From d2c70bd8b87beb54da40cd951ea5ba4ef7cfa463 Mon Sep 17 00:00:00 2001 From: Vladimir Iakovlev Date: Sun, 8 Oct 2017 16:17:41 +0200 Subject: [PATCH] #682: Use our own shell logger, fix experimental instant mode on macos --- tests/test_argument_parser.py | 8 +++- thefuck/argument_parser.py | 4 ++ thefuck/entrypoints/main.py | 3 ++ thefuck/entrypoints/shell_logger.py | 70 +++++++++++++++++++++++++++++ thefuck/shells/bash.py | 2 +- thefuck/shells/zsh.py | 2 +- 6 files changed, 85 insertions(+), 4 deletions(-) create mode 100644 thefuck/entrypoints/shell_logger.py diff --git a/tests/test_argument_parser.py b/tests/test_argument_parser.py index d2eb7c27..e5f1f44b 100644 --- a/tests/test_argument_parser.py +++ b/tests/test_argument_parser.py @@ -7,7 +7,8 @@ def _args(**override): args = {'alias': None, 'command': [], 'yes': False, 'help': False, 'version': False, 'debug': False, 'force_command': None, 'repeat': False, - 'enable_experimental_instant_mode': False} + 'enable_experimental_instant_mode': False, + 'shell_logger': None} args.update(override) return args @@ -27,6 +28,9 @@ def _args(**override): (['thefuck', 'git', 'branch', '-a', ARGUMENT_PLACEHOLDER, '-y', '-d'], _args(command=['git', 'branch', '-a'], yes=True, debug=True)), (['thefuck', 'git', 'branch', '-a', ARGUMENT_PLACEHOLDER, '-r', '-d'], - _args(command=['git', 'branch', '-a'], repeat=True, debug=True))]) + _args(command=['git', 'branch', '-a'], repeat=True, debug=True)), + (['thefuck', '-l', '/tmp/log'], _args(shell_logger='/tmp/log')), + (['thefuck', '--shell-logger', '/tmp/log'], + _args(shell_logger='/tmp/log'))]) def test_parse(argv, result): assert vars(Parser().parse(argv)) == result diff --git a/thefuck/argument_parser.py b/thefuck/argument_parser.py index 695111c1..8e79fc0a 100644 --- a/thefuck/argument_parser.py +++ b/thefuck/argument_parser.py @@ -25,6 +25,10 @@ class Parser(object): nargs='?', const=get_alias(), help='[custom-alias-name] prints alias for current shell') + self._parser.add_argument( + '-l', '--shell-logger', + action='store', + help='log shell output to the file') self._parser.add_argument( '--enable-experimental-instant-mode', action='store_true', diff --git a/thefuck/entrypoints/main.py b/thefuck/entrypoints/main.py index 587c7716..f41f0c3d 100644 --- a/thefuck/entrypoints/main.py +++ b/thefuck/entrypoints/main.py @@ -10,6 +10,7 @@ from ..argument_parser import Parser # noqa: E402 from ..utils import get_installation_info # noqa: E402 from .alias import print_alias # noqa: E402 from .fix_command import fix_command # noqa: E402 +from .shell_logger import shell_logger # noqa: E402 def main(): @@ -25,5 +26,7 @@ def main(): fix_command(known_args) elif known_args.alias: print_alias(known_args) + elif known_args.shell_logger: + shell_logger(known_args.shell_logger) else: parser.print_usage() diff --git a/thefuck/entrypoints/shell_logger.py b/thefuck/entrypoints/shell_logger.py new file mode 100644 index 00000000..6c903656 --- /dev/null +++ b/thefuck/entrypoints/shell_logger.py @@ -0,0 +1,70 @@ +import array +import fcntl +from functools import partial +import os +import pty +import signal +import sys +import termios +import tty +from ..logs import warn + + +def _read(f, fd): + data = os.read(fd, 1024) + f.write(data) + f.flush() + return data + + +def _set_pty_size(master_fd): + buf = array.array('h', [0, 0, 0, 0]) + fcntl.ioctl(pty.STDOUT_FILENO, termios.TIOCGWINSZ, buf, True) + fcntl.ioctl(master_fd, termios.TIOCSWINSZ, buf) + + +def _spawn(shell, master_read): + """Create a spawned process. + + Modified version of pty.spawn with terminal size support. + + """ + pid, master_fd = pty.fork() + + if pid == pty.CHILD: + os.execlp(shell, shell) + + try: + mode = tty.tcgetattr(pty.STDIN_FILENO) + tty.setraw(pty.STDIN_FILENO) + restore = True + except tty.error: # This is the same as termios.error + restore = False + + _set_pty_size(master_fd) + signal.signal(signal.SIGWINCH, lambda *_: _set_pty_size(master_fd)) + + try: + pty._copy(master_fd, master_read, pty._read) + except OSError: + if restore: + tty.tcsetattr(pty.STDIN_FILENO, tty.TCSAFLUSH, mode) + + os.close(master_fd) + return os.waitpid(pid, 0)[1] + + +def shell_logger(output): + """Logs shell output to the `output`. + + Works like unix script command with `-f` flag. + + """ + if not os.environ.get('SHELL'): + warn("Shell logger doesn't support your platform.") + sys.exit(1) + + with open(output, 'wb') as f: + return_code = _spawn(os.environ['SHELL'], partial(_read, f)) + + sys.exit(return_code) diff --git a/thefuck/shells/bash.py b/thefuck/shells/bash.py index 6eeaccab..724dffb0 100644 --- a/thefuck/shells/bash.py +++ b/thefuck/shells/bash.py @@ -44,7 +44,7 @@ class Bash(Generic): return ''' export THEFUCK_INSTANT_MODE=True; export THEFUCK_OUTPUT_LOG={log}; - script -feq {log}; + thefuck --shell-logger {log}; rm {log}; exit '''.format(log=log_path) diff --git a/thefuck/shells/zsh.py b/thefuck/shells/zsh.py index 6251076f..b134833c 100644 --- a/thefuck/shells/zsh.py +++ b/thefuck/shells/zsh.py @@ -42,7 +42,7 @@ class Zsh(Generic): return ''' export THEFUCK_INSTANT_MODE=True; export THEFUCK_OUTPUT_LOG={log}; - script -feq {log}; + thefuck --shell-logger {log}; rm {log}; exit '''.format(log=log_path)