#277 Fix the apt_get rule with sudo

tests/rules/test_apt_get.py

@@ -16,6 +16,8 @@ def test_match(command):
 
 @pytest.mark.parametrize('command, return_value', [
     (Command(script='vim', stderr='vim: command not found'),
+     [('vim', 'main'), ('vim-tiny', 'main')]),
+    (Command(script='sudo vim', stderr='vim: command not found'),
      [('vim', 'main'), ('vim-tiny', 'main')])])
 @patch('thefuck.rules.apt_get.CommandNotFound', create=True)
 @patch.multiple(apt_get, create=True, apt_get='apt_get')
@@ -38,7 +40,9 @@ def test_not_match(command):
                     reason='Skip if python-commandnotfound is not available')
 @pytest.mark.parametrize('command, new_command', [
     (Command('vim'), 'sudo apt-get install vim && vim'),
-    (Command('convert'), 'sudo apt-get install imagemagick && convert')])
+    (Command('convert'), 'sudo apt-get install imagemagick && convert'),
+    (Command('sudo vim'), 'sudo apt-get install vim && sudo vim'),
+    (Command('sudo convert'), 'sudo apt-get install imagemagick && sudo convert')])
 def test_get_new_command(command, new_command):
     assert get_new_command(command, None) == new_command
 
@@ -47,6 +51,11 @@ def test_get_new_command(command, new_command):
     (Command('vim'), 'sudo apt-get install vim && vim',
      [('vim', 'main'), ('vim-tiny', 'main')]),
     (Command('convert'), 'sudo apt-get install imagemagick && convert',
+     [('imagemagick', 'main'),
+      ('graphicsmagick-imagemagick-compat', 'universe')]),
+    (Command('sudo vim'), 'sudo apt-get install vim && sudo vim',
+     [('vim', 'main'), ('vim-tiny', 'main')]),
+    (Command('sudo convert'), 'sudo apt-get install imagemagick && sudo convert',
      [('imagemagick', 'main'),
       ('graphicsmagick-imagemagick-compat', 'universe')])])
 @patch('thefuck.rules.apt_get.CommandNotFound', create=True)
@@ -55,5 +64,3 @@ def test_get_new_command_mocked(cmdnf_mock, command, new_command, return_value):
     get_packages = Mock(return_value=return_value)
     cmdnf_mock.CommandNotFound.return_value = Mock(getPackages=get_packages)
     assert get_new_command(command, None) == new_command
-    assert cmdnf_mock.CommandNotFound.called
-    assert get_packages.called

thefuck/rules/apt_get.py

@@ -1,27 +1,30 @@
 from thefuck import shells
-from thefuck.utils import sudo_support
+from thefuck.utils import memoize
 
 try:
     import CommandNotFound
 except ImportError:
     enabled_by_default = False
 
-@sudo_support
-def match(command, settings):
-    if 'not found' in command.stderr:
-        try:
-            c = CommandNotFound.CommandNotFound()
-            pkgs = c.getPackages(command.script.split(" ")[0])
-            name, _ = pkgs[0]
-            return True
-        except IndexError:
-            # IndexError is thrown when no matching package is found
-            return False
 
-@sudo_support
+@memoize
+def get_package(command):
+    try:
+        c = CommandNotFound.CommandNotFound()
+        cmd = command.split(' ')
+        pkgs = c.getPackages(cmd[0] if cmd[0] != 'sudo' else cmd[1])
+        name, _ = pkgs[0]
+        return name
+    except IndexError:
+        # IndexError is thrown when no matching package is found
+        return None
+
+
+def match(command, settings):
+    return 'not found' in command.stderr and get_package(command.script)
+
+
 def get_new_command(command, settings):
-    c = CommandNotFound.CommandNotFound()
+    name = get_package(command.script)
-    pkgs = c.getPackages(command.script.split(" ")[0])
-    name, _ = pkgs[0]
     formatme = shells.and_('sudo apt-get install {}', '{}')
     return formatme.format(name, command.script)