Merge pull request #324 from nvbn/298-variants

Add ability to select fixed command from variants

thefuck/corrector.py

@@ -0,0 +1,80 @@
+import sys
+from imp import load_source
+from pathlib import Path
+from . import conf, types, logs
+from .utils import eager
+
+
+def load_rule(rule, settings):
+    """Imports rule module and returns it."""
+    name = rule.name[:-3]
+    rule_module = load_source(name, str(rule))
+    priority = getattr(rule_module, 'priority', conf.DEFAULT_PRIORITY)
+    return types.Rule(name, rule_module.match,
+                      rule_module.get_new_command,
+                      getattr(rule_module, 'enabled_by_default', True),
+                      getattr(rule_module, 'side_effect', None),
+                      settings.priority.get(name, priority),
+                      getattr(rule_module, 'requires_output', True))
+
+
+def get_loaded_rules(rules, settings):
+    """Yields all available rules."""
+    for rule in rules:
+        if rule.name != '__init__.py':
+            loaded_rule = load_rule(rule, settings)
+            if loaded_rule in settings.rules:
+                yield loaded_rule
+
+
+@eager
+def get_rules(user_dir, settings):
+    """Returns all enabled rules."""
+    bundled = Path(__file__).parent \
+        .joinpath('rules') \
+        .glob('*.py')
+    user = user_dir.joinpath('rules').glob('*.py')
+    return get_loaded_rules(sorted(bundled) + sorted(user), settings)
+
+
+@eager
+def get_matched_rules(command, rules, settings):
+    """Returns first matched rule for command."""
+    script_only = command.stdout is None and command.stderr is None
+
+    for rule in rules:
+        if script_only and rule.requires_output:
+            continue
+
+        try:
+            with logs.debug_time(u'Trying rule: {};'.format(rule.name),
+                                 settings):
+                if rule.match(command, settings):
+                    yield rule
+        except Exception:
+            logs.rule_failed(rule, sys.exc_info(), settings)
+
+
+def make_corrected_commands(command, rules, settings):
+    for rule in rules:
+        new_commands = rule.get_new_command(command, settings)
+        if not isinstance(new_commands, list):
+            new_commands = [new_commands]
+        for n, new_command in enumerate(new_commands):
+            yield types.CorrectedCommand(script=new_command,
+                                         side_effect=rule.side_effect,
+                                         priority=(n + 1) * rule.priority)
+
+
+def get_corrected_commands(command, user_dir, settings):
+    rules = get_rules(user_dir, settings)
+    logs.debug(
+        u'Loaded rules: {}'.format(', '.join(rule.name for rule in rules)),
+        settings)
+    matched = get_matched_rules(command, rules, settings)
+    logs.debug(
+        u'Matched rules: {}'.format(', '.join(rule.name for rule in matched)),
+        settings)
+    corrected_commands = make_corrected_commands(command, matched, settings)
+    return sorted(corrected_commands,
+                  key=lambda corrected_command: corrected_command.priority)

thefuck/logs.py

@@ -1,3 +1,5 @@
+# -*- encoding: utf-8 -*-
+
 from contextlib import contextmanager
 from datetime import datetime
 import sys
@@ -28,27 +30,6 @@ def rule_failed(rule, exc_info, settings):
     exception('Rule {}'.format(rule.name), exc_info, settings)
 
 
-def show_command(new_command, side_effect, settings):
-    sys.stderr.write('{bold}{command}{reset}{side_effect}\n'.format(
-        command=new_command,
-        side_effect=' (+side effect)' if side_effect else '',
-        bold=color(colorama.Style.BRIGHT, settings),
-        reset=color(colorama.Style.RESET_ALL, settings)))
-
-
-def confirm_command(new_command, side_effect, settings):
-    sys.stderr.write(
-        '{bold}{command}{reset}{side_effect} '
-        '[{green}enter{reset}/{red}ctrl+c{reset}]'.format(
-            command=new_command,
-            side_effect=' (+side effect)' if side_effect else '',
-            bold=color(colorama.Style.BRIGHT, settings),
-            green=color(colorama.Fore.GREEN, settings),
-            red=color(colorama.Fore.RED, settings),
-            reset=color(colorama.Style.RESET_ALL, settings)))
-    sys.stderr.flush()
-
-
 def failed(msg, settings):
     sys.stderr.write('{red}{msg}{reset}\n'.format(
         msg=msg,
@@ -56,6 +37,27 @@ def failed(msg, settings):
         reset=color(colorama.Style.RESET_ALL, settings)))
 
 
+def show_corrected_command(corrected_command, settings):
+    sys.stderr.write('{bold}{script}{reset}{side_effect}\n'.format(
+        script=corrected_command.script,
+        side_effect=' (+side effect)' if corrected_command.side_effect else '',
+        bold=color(colorama.Style.BRIGHT, settings),
+        reset=color(colorama.Style.RESET_ALL, settings)))
+
+
+def confirm_text(corrected_command, settings):
+    sys.stderr.write(
+        '\033[1K\r{bold}{script}{reset}{side_effect} '
+        '[{green}enter{reset}/{blue}↑{reset}/{blue}↓{reset}/{red}ctrl+c{reset}]'.format(
+            script=corrected_command.script,
+            side_effect=' (+side effect)' if corrected_command.side_effect else '',
+            bold=color(colorama.Style.BRIGHT, settings),
+            green=color(colorama.Fore.GREEN, settings),
+            red=color(colorama.Fore.RED, settings),
+            reset=color(colorama.Style.RESET_ALL, settings),
+            blue=color(colorama.Fore.BLUE, settings)))
+
+
 def debug(msg, settings):
     if settings.debug:
         sys.stderr.write(u'{blue}{bold}DEBUG:{reset} {msg}\n'.format(

thefuck/main.py

@@ -1,4 +1,3 @@
-from imp import load_source
 from pathlib import Path
 from os.path import expanduser
 from pprint import pformat
@@ -9,6 +8,8 @@ from psutil import Process, TimeoutExpired
 import colorama
 import six
 from . import logs, conf, types, shells
+from .corrector import get_corrected_commands
+from .ui import select_command
 
 
 def setup_user_dir():
@@ -21,37 +22,6 @@ def setup_user_dir():
     return user_dir
 
 
-def load_rule(rule):
-    """Imports rule module and returns it."""
-    rule_module = load_source(rule.name[:-3], str(rule))
-    return types.Rule(rule.name[:-3], rule_module.match,
-                      rule_module.get_new_command,
-                      getattr(rule_module, 'enabled_by_default', True),
-                      getattr(rule_module, 'side_effect', None),
-                      getattr(rule_module, 'priority', conf.DEFAULT_PRIORITY),
-                      getattr(rule_module, 'requires_output', True))
-
-
-def _get_loaded_rules(rules, settings):
-    """Yields all available rules."""
-    for rule in rules:
-        if rule.name != '__init__.py':
-            loaded_rule = load_rule(rule)
-            if loaded_rule in settings.rules:
-                yield loaded_rule
-
-
-def get_rules(user_dir, settings):
-    """Returns all enabled rules."""
-    bundled = Path(__file__).parent \
-        .joinpath('rules') \
-        .glob('*.py')
-    user = user_dir.joinpath('rules').glob('*.py')
-    rules = _get_loaded_rules(sorted(bundled) + sorted(user), settings)
-    return sorted(rules, key=lambda rule: settings.priority.get(
-        rule.name, rule.priority))
-
-
 def wait_output(settings, popen):
     """Returns `True` if we can get output of the command in the
     `wait_command` time.
@@ -100,46 +70,12 @@ def get_command(settings, args):
             return types.Command(script, None, None)
 
 
-def get_matched_rule(command, rules, settings):
-    """Returns first matched rule for command."""
-    script_only = command.stdout is None and command.stderr is None
-
-    for rule in rules:
-        if script_only and rule.requires_output:
-            continue
-
-        try:
-            with logs.debug_time(u'Trying rule: {};'.format(rule.name),
-                                 settings):
-                if rule.match(command, settings):
-                    return rule
-        except Exception:
-            logs.rule_failed(rule, sys.exc_info(), settings)
-
-
-def confirm(new_command, side_effect, settings):
-    """Returns `True` when running of new command confirmed."""
-    if not settings.require_confirmation:
-        logs.show_command(new_command, side_effect, settings)
-        return True
-
-    logs.confirm_command(new_command, side_effect, settings)
-    try:
-        sys.stdin.read(1)
-        return True
-    except KeyboardInterrupt:
-        logs.failed('Aborted', settings)
-        return False
-
-
-def run_rule(rule, command, settings):
+def run_command(command, settings):
     """Runs command from rule for passed command."""
-    new_command = shells.to_shell(rule.get_new_command(command, settings))
-    if confirm(new_command, rule.side_effect, settings):
-        if rule.side_effect:
-            rule.side_effect(command, settings)
-        shells.put_to_history(new_command)
-        print(new_command)
+    if command.side_effect:
+        command.side_effect(command, settings)
+    shells.put_to_history(command.script)
+    print(command.script)
 
 
 # Entry points:
@@ -152,18 +88,10 @@ def main():
         logs.debug(u'Run with settings: {}'.format(pformat(settings)), settings)
 
         command = get_command(settings, sys.argv)
-        rules = get_rules(user_dir, settings)
-        logs.debug(
-            u'Loaded rules: {}'.format(', '.join(rule.name for rule in rules)),
-            settings)
-
-        matched_rule = get_matched_rule(command, rules, settings)
-        if matched_rule:
-            logs.debug(u'Matched rule: {}'.format(matched_rule.name), settings)
-            run_rule(matched_rule, command, settings)
-            return
-
-        logs.failed('No fuck given', settings)
+        corrected_commands = get_corrected_commands(command, user_dir, settings)
+        selected_command = select_command(corrected_commands, settings)
+        if selected_command:
+            run_command(selected_command, settings)
 
 
 def print_alias():

thefuck/rules/brew_unknown_command.py

@@ -1,7 +1,7 @@
 import os
 import re
 import subprocess
-from thefuck.utils import get_closest, replace_argument
+from thefuck.utils import get_closest, replace_command
 
 BREW_CMD_PATH = '/Library/Homebrew/cmd'
 TAP_PATH = '/Library/Taps'
@@ -77,10 +77,6 @@ if brew_path_prefix:
         pass
 
 
-def _get_similar_command(command):
-    return get_closest(command, brew_commands)
-
-
 def match(command, settings):
     is_proper_command = ('brew' in command.script and
                          'Unknown command' in command.stderr)
@@ -89,7 +85,7 @@ def match(command, settings):
     if is_proper_command:
         broken_cmd = re.findall(r'Error: Unknown command: ([a-z]+)',
                                 command.stderr)[0]
-        has_possible_commands = bool(_get_similar_command(broken_cmd))
+        has_possible_commands = bool(get_closest(broken_cmd, brew_commands))
 
     return has_possible_commands
 
@@ -97,6 +93,4 @@ def match(command, settings):
 def get_new_command(command, settings):
     broken_cmd = re.findall(r'Error: Unknown command: ([a-z]+)',
                             command.stderr)[0]
-    new_cmd = _get_similar_command(broken_cmd)
-
-    return replace_argument(command.script, broken_cmd, new_cmd)
+    return replace_command(command, broken_cmd, brew_commands)

thefuck/rules/docker_not_command.py

@@ -1,7 +1,7 @@
 from itertools import dropwhile, takewhile, islice
 import re
 import subprocess
-from thefuck.utils import get_closest, sudo_support, replace_argument
+from thefuck.utils import get_closest, sudo_support, replace_argument, replace_command
 
 
 @sudo_support
@@ -23,5 +23,4 @@ def get_docker_commands():
 def get_new_command(command, settings):
     wrong_command = re.findall(
         r"docker: '(\w+)' is not a docker command.", command.stderr)[0]
-    fixed_command = get_closest(wrong_command, get_docker_commands())
-    return replace_argument(command.script, wrong_command, fixed_command)
+    return replace_command(command, wrong_command, get_docker_commands())

thefuck/rules/git_not_command.py

@@ -1,6 +1,6 @@
 import re
-from thefuck.utils import (get_closest, git_support, replace_argument,
-                           get_all_matched_commands)
+from thefuck.utils import (git_support,
+                           get_all_matched_commands, replace_command)
 
 
 @git_support
@@ -13,7 +13,5 @@ def match(command, settings):
 def get_new_command(command, settings):
     broken_cmd = re.findall(r"git: '([^']*)' is not a git command",
                             command.stderr)[0]
-    new_cmd = get_closest(broken_cmd,
-                          get_all_matched_commands(command.stderr))
-    return replace_argument(command.script, broken_cmd, new_cmd)
-
+    matched = get_all_matched_commands(command.stderr)
+    return replace_command(command, broken_cmd, matched)

thefuck/rules/gulp_not_task.py

@@ -1,6 +1,6 @@
 import re
 import subprocess
-from thefuck.utils import get_closest, replace_argument
+from thefuck.utils import replace_command
 
 
 def match(command, script):
@@ -18,5 +18,4 @@ def get_gulp_tasks():
 def get_new_command(command, script):
     wrong_task = re.findall(r"Task '(\w+)' is not in your gulpfile",
                             command.stdout)[0]
-    fixed_task = get_closest(wrong_task, get_gulp_tasks())
-    return replace_argument(command.script, wrong_task, fixed_task)
+    return replace_command(command, wrong_task, get_gulp_tasks())

thefuck/rules/heroku_not_command.py

@@ -1,5 +1,5 @@
 import re
-from thefuck.utils import get_closest, replace_argument
+from thefuck.utils import replace_command
 
 
 def match(command, settings):
@@ -16,5 +16,4 @@ def _get_suggests(stderr):
 
 def get_new_command(command, settings):
     wrong = re.findall(r'`(\w+)` is not a heroku command', command.stderr)[0]
-    correct = get_closest(wrong, _get_suggests(command.stderr))
-    return replace_argument(command.script, wrong, correct)
+    return replace_command(command, wrong, _get_suggests(command.stderr))

thefuck/rules/lein_not_task.py

@@ -1,5 +1,6 @@
 import re
-from thefuck.utils import sudo_support, replace_argument
+from thefuck.utils import sudo_support,\
+    replace_command, get_all_matched_commands
 
 
 @sudo_support
@@ -13,6 +14,5 @@ def match(command, settings):
 def get_new_command(command, settings):
     broken_cmd = re.findall(r"'([^']*)' is not a task",
                             command.stderr)[0]
-    new_cmd = re.findall(r'Did you mean this\?\n\s*([^\n]*)',
-                         command.stderr)[0]
-    return replace_argument(command.script, broken_cmd, new_cmd)
+    new_cmds = get_all_matched_commands(command.stderr, 'Did you mean this?')
+    return replace_command(command, broken_cmd, new_cmds)

thefuck/rules/no_command.py

@@ -1,5 +1,5 @@
 from difflib import get_close_matches
-from thefuck.utils import sudo_support, get_all_executables, get_closest
+from thefuck.utils import sudo_support, get_all_executables
 
 
 @sudo_support
@@ -12,8 +12,9 @@ def match(command, settings):
 @sudo_support
 def get_new_command(command, settings):
     old_command = command.script.split(' ')[0]
-    new_command = get_closest(old_command, get_all_executables())
-    return ' '.join([new_command] + command.script.split(' ')[1:])
+    new_cmds = get_close_matches(old_command, get_all_executables(), cutoff=0.1)
+    return [' '.join([new_command] + command.script.split(' ')[1:])
+            for new_command in new_cmds]
 
 
 priority = 3000

thefuck/rules/tsuru_not_command.py

@@ -1,6 +1,6 @@
 import re
 from thefuck.utils import (get_closest, replace_argument,
-                           get_all_matched_commands)
+                           get_all_matched_commands, replace_command)
 
 
 def match(command, settings):
@@ -12,7 +12,6 @@ def match(command, settings):
 def get_new_command(command, settings):
     broken_cmd = re.findall(r'tsuru: "([^"]*)" is not a tsuru command',
                             command.stderr)[0]
-    new_cmd = get_closest(broken_cmd,
-                          get_all_matched_commands(command.stderr))
-    return replace_argument(command.script, broken_cmd, new_cmd)
+    return replace_command(command, broken_cmd,
+                           get_all_matched_commands(command.stderr))
 

thefuck/types.py

@@ -3,6 +3,8 @@ from collections import namedtuple
 
 Command = namedtuple('Command', ('script', 'stdout', 'stderr'))
 
+CorrectedCommand = namedtuple('CorrectedCommand', ('script', 'side_effect', 'priority'))
+
 Rule = namedtuple('Rule', ('name', 'match', 'get_new_command',
                            'enabled_by_default', 'side_effect',
                            'priority', 'requires_output'))

thefuck/ui.py

@@ -0,0 +1,103 @@
+# -*- encoding: utf-8 -*-
+
+import sys
+from . import logs
+
+try:
+    from msvcrt import getch
+except ImportError:
+    def getch():
+        import tty
+        import termios
+
+        fd = sys.stdin.fileno()
+        old = termios.tcgetattr(fd)
+        try:
+            tty.setraw(fd)
+            ch = sys.stdin.read(1)
+            if ch == '\x03':  # For compatibility with msvcrt.getch
+                raise KeyboardInterrupt
+            return ch
+        finally:
+            termios.tcsetattr(fd, termios.TCSADRAIN, old)
+
+SELECT = 0
+ABORT = 1
+PREVIOUS = 2
+NEXT = 3
+
+
+def read_actions():
+    """Yields actions for pressed keys."""
+    buffer = []
+    while True:
+        try:
+            ch = getch()
+        except KeyboardInterrupt:  # Ctrl+C
+            yield ABORT
+
+        if ch in ('\n', '\r'):  # Enter
+            yield SELECT
+
+        buffer.append(ch)
+        buffer = buffer[-3:]
+
+        if buffer == ['\x1b', '[', 'A']:  # ↑
+            yield PREVIOUS
+
+        if buffer == ['\x1b', '[', 'B']:  # ↓
+            yield NEXT
+
+
+class CommandSelector(object):
+    def __init__(self, commands):
+        self._commands = commands
+        self._index = 0
+        self._on_change = lambda x: x
+
+    def next(self):
+        self._index = (self._index + 1) % len(self._commands)
+        self._on_change(self.value)
+
+    def previous(self):
+        self._index = (self._index - 1) % len(self._commands)
+        self._on_change(self.value)
+
+    @property
+    def value(self):
+        return self._commands[self._index]
+
+    def on_change(self, fn):
+        self._on_change = fn
+        fn(self.value)
+
+
+def select_command(corrected_commands, settings):
+    """Returns:
+
+     - the first command when confirmation disabled;
+     - None when ctrl+c pressed;
+     - selected command.
+
+    """
+    if not corrected_commands:
+        logs.failed('No fuck given', settings)
+        return
+
+    selector = CommandSelector(corrected_commands)
+    if not settings.require_confirmation:
+        logs.show_corrected_command(selector.value, settings)
+        return selector.value
+
+    selector.on_change(lambda val: logs.confirm_text(val, settings))
+    for action in read_actions():
+        if action == SELECT:
+            sys.stderr.write('\n')
+            return selector.value
+        elif action == ABORT:
+            logs.failed('\nAborted', settings)
+            return
+        elif action == PREVIOUS:
+            selector.previous()
+        elif action == NEXT:
+            selector.next()

thefuck/utils.py

@@ -69,6 +69,8 @@ def sudo_support(fn):
 
         if result and isinstance(result, six.string_types):
             return u'sudo {}'.format(result)
+        elif isinstance(result, list):
+            return [u'sudo {}'.format(x) for x in result]
         else:
             return result
     return wrapper
@@ -161,6 +163,14 @@ def replace_argument(script, from_, to):
             u' {} '.format(from_), u' {} '.format(to), 1)
 
 
+def eager(fn):
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        return list(fn(*args, **kwargs))
+    return wrapper
+
+
+@eager
 def get_all_matched_commands(stderr, separator='Did you mean'):
     should_yield = False
     for line in stderr.split('\n'):
@@ -168,3 +178,10 @@ def get_all_matched_commands(stderr, separator='Did you mean'):
             should_yield = True
         elif should_yield and line:
             yield line.strip()
+
+
+def replace_command(command, broken, matched):
+    """Helper for *_no_command rules."""
+    new_cmds = get_close_matches(broken, matched, cutoff=0.1)
+    return [replace_argument(command.script, broken, new_cmd.strip())
+            for new_cmd in new_cmds]