#21 Add timeout for getting previous command output

README.md

@@ -122,6 +122,7 @@ def get_new_command(command, settings):
 The Fuck have a few settings parameters:
 
 * `rules` – list of enabled rules, by default all;
+* `wait_command` – max amount of time in seconds for getting previous command output;
 * `command_not_found` – path to `command_not_found` binary,
 by default `/usr/lib/command-not-found`.
 

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 
 setup(name='thefuck',
-      version=1.10,
+      version=1.11,
       description="Magnificent app which corrects your previous console command",
       author='Vladimir Iakovlev',
       author_email='nvbn.rm@gmail.com',
@@ -11,6 +11,6 @@ setup(name='thefuck',
       packages=find_packages(exclude=['ez_setup', 'examples', 'tests']),
       include_package_data=True,
       zip_safe=False,
-      install_requires=['pathlib'],
+      install_requires=['pathlib', 'psutil'],
       entry_points={'console_scripts': [
           'thefuck = thefuck.main:main']})

tests/test_main.py

@@ -36,22 +36,25 @@ def test_get_rules():
         assert main.get_rules(
             Path('~'),
             Mock(rules=None)) == [main.Rule('bash', 'bash'),
-                                     main.Rule('lisp', 'lisp'),
-                                     main.Rule('bash', 'bash'),
-                                     main.Rule('lisp', 'lisp')]
+                                  main.Rule('lisp', 'lisp'),
+                                  main.Rule('bash', 'bash'),
+                                  main.Rule('lisp', 'lisp')]
         assert main.get_rules(
             Path('~'),
             Mock(rules=['bash'])) == [main.Rule('bash', 'bash'),
-                                         main.Rule('bash', 'bash')]
+                                      main.Rule('bash', 'bash')]
 
 
 def test_get_command():
-    with patch('thefuck.main.Popen') as Popen,\
+    with patch('thefuck.main.Popen') as Popen, \
             patch('thefuck.main.os.environ',
-                  new_callable=lambda: {}):
+                  new_callable=lambda: {}), \
+            patch('thefuck.main.wait_output',
+                  return_value=True):
         Popen.return_value.stdout.read.return_value = b'stdout'
         Popen.return_value.stderr.read.return_value = b'stderr'
-        assert main.get_command(['thefuck', 'apt-get', 'search', 'vim']) \
+        assert main.get_command(Mock(), ['thefuck', 'apt-get',
+                                         'search', 'vim']) \
                == main.Command('apt-get search vim', 'stdout', 'stderr')
         Popen.assert_called_once_with('apt-get search vim',
                                       shell=True,

thefuck/main.py

@@ -5,6 +5,7 @@ from os.path import expanduser
 from subprocess import Popen, PIPE
 import os
 import sys
+from psutil import Process, TimeoutExpired
 
 
 Command = namedtuple('Command', ('script', 'stdout', 'stderr'))
@@ -25,8 +26,8 @@ def get_settings(user_dir):
     """Returns prepared settings module."""
     settings = load_source('settings',
                            str(user_dir.joinpath('settings.py')))
-    if not hasattr(settings, 'rules'):
-        settings.rules = None
+    settings.__dict__.setdefault('rules', None)
+    settings.__dict__.setdefault('wait_command', 3)
     return settings
 
 
@@ -54,13 +55,32 @@ def get_rules(user_dir, settings):
             if rule.name != '__init__.py' and is_rule_enabled(settings, rule)]
 
 
-def get_command(args):
+def wait_output(settings, popen):
+    """Returns `True` if we can get output of the command in the
+    `wait_command` time.
+
+    Command will be killed if it wasn't finished in the time.
+
+    """
+    proc = Process(popen.pid)
+    try:
+        proc.wait(settings.wait_command)
+        return True
+    except TimeoutExpired:
+        for child in proc.get_children(recursive=True):
+            child.kill()
+        proc.kill()
+        return False
+
+
+def get_command(settings, args):
     """Creates command from `args` and executes it."""
     script = ' '.join(args[1:])
     result = Popen(script, shell=True, stdout=PIPE, stderr=PIPE,
                    env=dict(os.environ, LANG='C'))
-    return Command(script, result.stdout.read().decode('utf-8'),
-                   result.stderr.read().decode('utf-8'))
+    if wait_output(settings, result):
+        return Command(script, result.stdout.read().decode('utf-8'),
+                       result.stderr.read().decode('utf-8'))
 
 
 def get_matched_rule(command, rules, settings):
@@ -83,15 +103,19 @@ def is_second_run(command):
 
 
 def main():
-    command = get_command(sys.argv)
-    if is_second_run(command):
-        print("echo Can't fuck twice")
-    else:
-        user_dir = setup_user_dir()
-        settings = get_settings(user_dir)
+    user_dir = setup_user_dir()
+    settings = get_settings(user_dir)
+
+    command = get_command(settings, sys.argv)
+    if command:
+        if is_second_run(command):
+            print("echo Can't fuck twice")
+            return
+
         rules = get_rules(user_dir, settings)
         matched_rule = get_matched_rule(command, rules, settings)
         if matched_rule:
             run_rule(matched_rule, command, settings)
-        else:
-            print('echo No fuck given')
+            return
+
+    print('echo No fuck given')