diff --git a/setup.py b/setup.py index 701e9af..b9470cc 100644 --- a/setup.py +++ b/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']}) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..fed491a --- /dev/null +++ b/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' diff --git a/thefuck/rules/cp_omitting_directory.py b/thefuck/rules/cp_omitting_directory.py index 14e84c9..4dee781 100644 --- a/thefuck/rules/cp_omitting_directory.py +++ b/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) diff --git a/thefuck/rules/has_exists_script.py b/thefuck/rules/has_exists_script.py index 4ceac48..19a7e48 100644 --- a/thefuck/rules/has_exists_script.py +++ b/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) diff --git a/thefuck/rules/lein_not_task.py b/thefuck/rules/lein_not_task.py index efc25a1..a043263 100644 --- a/thefuck/rules/lein_not_task.py +++ b/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] diff --git a/thefuck/rules/mkdir_p.py b/thefuck/rules/mkdir_p.py index 896f08f..03b40ce 100644 --- a/thefuck/rules/mkdir_p.py +++ b/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) diff --git a/thefuck/rules/no_command.py b/thefuck/rules/no_command.py index ce90926..917e1fc 100644 --- a/thefuck/rules/no_command.py +++ b/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, diff --git a/thefuck/rules/python_command.py b/thefuck/rules/python_command.py index 507a934..f2bc8dc 100644 --- a/thefuck/rules/python_command.py +++ b/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 diff --git a/thefuck/rules/rm_dir.py b/thefuck/rules/rm_dir.py index f9349ea..557278b 100644 --- a/thefuck/rules/rm_dir.py +++ b/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) diff --git a/thefuck/rules/rm_root.py b/thefuck/rules/rm_root.py index d07543b..ed0121f 100644 --- a/thefuck/rules/rm_root.py +++ b/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) diff --git a/thefuck/utils.py b/thefuck/utils.py index fa4ee1e..534a634 100644 --- a/thefuck/utils.py +++ b/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