From e745f3d4a95ad8c8479955a45e624c1081ff2c86 Mon Sep 17 00:00:00 2001 From: nvbn Date: Sat, 18 Apr 2015 22:50:18 +0200 Subject: [PATCH] #21 Add timeout for getting previous command output --- README.md | 1 + setup.py | 4 ++-- tests/test_main.py | 17 +++++++++------- thefuck/main.py | 50 ++++++++++++++++++++++++++++++++++------------ 4 files changed, 50 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index 9cdbfb6..e799f60 100644 --- a/README.md +++ b/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`. diff --git a/setup.py b/setup.py index b4cd609..5759cdc 100644 --- a/setup.py +++ b/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']}) diff --git a/tests/test_main.py b/tests/test_main.py index 8eb7605..84a9ebe 100644 --- a/tests/test_main.py +++ b/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, diff --git a/thefuck/main.py b/thefuck/main.py index e69046d..ae31919 100644 --- a/thefuck/main.py +++ b/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')