Add transparent sudo support for rules where it required

setup.py

@@ -15,6 +15,6 @@ setup(name='thefuck',
                                       'tests', 'release']),
       include_package_data=True,
       zip_safe=False,
-      install_requires=['pathlib', 'psutil', 'colorama'],
+      install_requires=['pathlib', 'psutil', 'colorama', 'six'],
       entry_points={'console_scripts': [
           'thefuck = thefuck.main:main']})

tests/test_utils.py

@@ -0,0 +1,17 @@
+from mock import Mock
+from thefuck.utils import sudo_support
+from thefuck.main import Command
+
+
+def test_sudo_support():
+    fn = Mock(return_value=True, __name__='')
+    assert sudo_support(fn)(Command('sudo ls', 'out', 'err'), None)
+    fn.assert_called_once_with(Command('ls', 'out', 'err'), None)
+
+    fn.return_value = False
+    assert not sudo_support(fn)(Command('sudo ls', 'out', 'err'), None)
+
+    fn.return_value = 'pwd'
+    assert sudo_support(fn)(Command('sudo ls', 'out', 'err'), None) == 'sudo pwd'
+
+    assert sudo_support(fn)(Command('ls', 'out', 'err'), None) == 'pwd'

thefuck/rules/cp_omitting_directory.py

@@ -1,10 +1,13 @@
 import re
+from thefuck.utils import sudo_support
 
 
+@sudo_support
 def match(command, settings):
     return command.script.startswith('cp ') \
         and 'cp: omitting directory' in command.stderr.lower()
 
 
+@sudo_support
 def get_new_command(command, settings):
     return re.sub(r'^cp', 'cp -a', command.script)

thefuck/rules/has_exists_script.py

@@ -1,11 +1,14 @@
 import os
+from thefuck.utils import sudo_support
 
 
+@sudo_support
 def match(command, settings):
     return os.path.exists(command.script.split()[0]) \
         and 'command not found' in command.stderr
 
 
+@sudo_support
 def get_new_command(command, settings):
     return u'./{}'.format(command.script)
 

thefuck/rules/lein_not_task.py

@@ -1,12 +1,15 @@
 import re
+from thefuck.utils import sudo_support
 
 
+@sudo_support
 def match(command, settings):
     return (command.script.startswith('lein')
             and "is not a task. See 'lein help'" in command.stderr
             and 'Did you mean this?' in command.stderr)
 
 
+@sudo_support
 def get_new_command(command, settings):
     broken_cmd = re.findall(r"'([^']*)' is not a task",
                             command.stderr)[0]

thefuck/rules/mkdir_p.py

@@ -1,9 +1,13 @@
 import re
+from thefuck.utils import sudo_support
 
+
+@sudo_support
 def match(command, settings):
     return ('mkdir' in command.script
             and 'No such file or directory' in command.stderr)
 
 
+@sudo_support
 def get_new_command(command, settings):
     return re.sub('^mkdir (.*)', 'mkdir -p \\1', command.script)

thefuck/rules/no_command.py

@@ -1,6 +1,7 @@
 from difflib import get_close_matches
 import os
 from pathlib import Path
+from thefuck.utils import sudo_support
 
 
 def _safe(fn, fallback):
@@ -17,12 +18,14 @@ def _get_all_bins():
             if not _safe(exe.is_dir, True)]
 
 
+@sudo_support
 def match(command, settings):
     return 'not found' in command.stderr and \
            bool(get_close_matches(command.script.split(' ')[0],
                                   _get_all_bins()))
 
 
+@sudo_support
 def get_new_command(command, settings):
     old_command = command.script.split(' ')[0]
     new_command = get_close_matches(old_command,

thefuck/rules/python_command.py

@@ -1,7 +1,10 @@
+from thefuck.utils import sudo_support
 # add 'python' suffix to the command if
 #  1) The script does not have execute permission or
 #  2) is interpreted as shell script
 
+
+@sudo_support
 def match(command, settings):
     toks = command.script.split()
     return (len(toks) > 0
@@ -10,5 +13,6 @@ def match(command, settings):
                  'command not found' in command.stderr))
 
 
+@sudo_support
 def get_new_command(command, settings):
     return 'python ' + command.script

thefuck/rules/rm_dir.py

@@ -1,9 +1,13 @@
 import re
+from thefuck.utils import sudo_support
 
+
+@sudo_support
 def match(command, settings):
     return ('rm' in command.script
             and 'is a directory' in command.stderr)
 
 
+@sudo_support
 def get_new_command(command, settings):
     return re.sub('^rm (.*)', 'rm -rf \\1', command.script)

thefuck/rules/rm_root.py

@@ -1,11 +1,16 @@
+from thefuck.utils import sudo_support
+
+
 enabled_by_default = False
 
 
+@sudo_support
 def match(command, settings):
     return ({'rm', '/'}.issubset(command.script.split())
             and '--no-preserve-root' not in command.script
             and '--no-preserve-root' in command.stderr)
 
 
+@sudo_support
 def get_new_command(command, settings):
     return u'{} --no-preserve-root'.format(command.script)

thefuck/utils.py

@@ -1,5 +1,7 @@
 from functools import wraps
 import os
+import six
+from thefuck.main import Command
 
 
 def which(program):
@@ -41,3 +43,22 @@ def wrap_settings(params):
             return fn(command, settings)
         return wrapper
     return decorator
+
+
+def sudo_support(fn):
+    """Removes sudo before calling fn and adds it after."""
+    @wraps(fn)
+    def wrapper(command, settings):
+        if not command.script.startswith('sudo '):
+            return fn(command, settings)
+
+        result = fn(Command(command.script[5:],
+                            command.stdout,
+                            command.stderr),
+                    settings)
+
+        if result and isinstance(result, six.string_types):
+            return u'sudo {}'.format(result)
+        else:
+            return result
+    return wrapper