├── .gitignore ├── Makefile ├── README.md ├── bin └── parent-lifetime ├── setup.py └── tee_output └── __init__.py /.gitignore: -------------------------------------------------------------------------------- 1 | /build 2 | /dist 3 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: build upload 2 | 3 | build: 4 | python setup.py sdist bdist_wheel 5 | 6 | upload: 7 | twine upload dist/* 8 | 9 | clean: 10 | rm -rf dist/ 11 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Tee Output 2 | 3 | `tee-output` is a Python library to tee standard output / standard 4 | error from the current process into a logfile. Unlike doing a shell 5 | redirection (i.e. `python myscript.py 2>&1 | tee /tmp/log`), 6 | `tee-output` preserves terminal semantics, so `breakpoint()` etc 7 | continue to work. 8 | 9 | Basic usage: 10 | 11 | ``` 12 | from tee_output import tee 13 | tee("/tmp/log.out", "/tmp/log.err") 14 | ``` 15 | 16 | After running the above, your standard output will stream to 17 | `/tmp/stdout.out` in addition to the terminal; your standard error will 18 | stream to `/tmp/stdout.err` in addition to the terminal. 19 | 20 | You can also provide a list of locations to `tee`: 21 | 22 | ``` 23 | from tee_output import tee 24 | tee(["/tmp/log.out", "/tmp/log.combined"], ["/tmp/log.err", "/tmp/log.combined"]) 25 | ``` 26 | 27 | This will additionally create a file `/tmp/log.combined` which 28 | contains the interleaved standard output and error, such as you see in 29 | your terminal. 30 | -------------------------------------------------------------------------------- /bin/parent-lifetime: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import os 3 | import psutil 4 | import subprocess 5 | import signal 6 | import sys 7 | import time 8 | 9 | poll_timeout = 1 10 | death_timeout = 10 11 | 12 | def close_fds(exclude={0, 1, 2}): 13 | if sys.platform == 'linux': 14 | fds = [int(fd) for fd in os.listdir("/proc/self/fd")] 15 | else: 16 | # Best effort guess 17 | fds = list(range(1024)) 18 | for fd in fds: 19 | if fd in exclude: 20 | continue 21 | try: 22 | os.close(fd) 23 | except OSError: 24 | pass 25 | 26 | # On Linux, could also use prctl to get a notification on parent death 27 | def main(): 28 | if len(sys.argv) >= 1 and sys.argv[1] == "--term": 29 | sys.argv = sys.argv[1:] 30 | term = True 31 | else: 32 | term = False 33 | 34 | def terminate(proc): 35 | if term: 36 | proc.terminate() 37 | else: 38 | proc.kill() 39 | sig = signal.SIGTERM if term else signal.SIGKILL 40 | 41 | if len(sys.argv) <= 1: 42 | sys.stderr.write(f"Usage: {sys.argv[0]} \n") 43 | return 1 44 | 45 | if sys.platform == 'linux': 46 | import ctypes 47 | 48 | def preexec_fn(): 49 | signal.signal(signal.SIGINT, signal.SIG_DFL) 50 | PR_SET_PDEATHSIG = 1 51 | ctypes.cdll["libc.so.6"].prctl(PR_SET_PDEATHSIG, int(sig)) 52 | else: 53 | def preexec_fn(): 54 | signal.signal(signal.SIGINT, signal.SIG_DFL) 55 | 56 | child = subprocess.Popen(sys.argv[1:], close_fds=False, preexec_fn=preexec_fn) 57 | close_fds(exclude={1, 2}) 58 | 59 | def sigint(sig, frame): 60 | terminate(child) 61 | child.wait() 62 | os._exit(0) 63 | 64 | signal.signal(signal.SIGINT, sigint) 65 | 66 | # Let the process die naturally 67 | while True: 68 | if child.poll() is not None: 69 | return child.returncode 70 | elif os.getppid() == 1: 71 | break 72 | time.sleep(poll_timeout) 73 | 74 | # Try to kill (and also carefully kill any descendents) 75 | proc = psutil.Process() 76 | descendents = proc.children(recursive=True) 77 | 78 | # Send TERM to our direct subprocess, giving it a chance to 79 | # cleanly exit. 80 | child.terminate() 81 | 82 | # Wait for it to die on its own, or time out and kill it directly. 83 | start = time.time() 84 | while time.time() - start < death_timeout: 85 | if child.poll() is not None: 86 | break 87 | 88 | # Now hard-kill any leftover processes 89 | for desc in descendents: 90 | try: 91 | desc.kill() 92 | except psutil.NoSuchProcess: 93 | pass 94 | 95 | child.wait() 96 | return child.returncode 97 | 98 | 99 | if __name__ == "__main__": 100 | sys.exit(main()) 101 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup(name='tee-output', 4 | version='0.4.16', 5 | description='A utility to tee standard output / standard error from the current process into a logfile. Preserves terminal semantics, so breakpoint() etc continue to work.', 6 | install_requires=['psutil'], 7 | author='Greg Brockman', 8 | packages=['tee_output'], 9 | scripts=["bin/parent-lifetime"] 10 | ) 11 | -------------------------------------------------------------------------------- /tee_output/__init__.py: -------------------------------------------------------------------------------- 1 | import fcntl 2 | import os 3 | import signal 4 | import struct 5 | import subprocess 6 | import sys 7 | import termios 8 | import tty 9 | from dataclasses import dataclass, field 10 | from typing import List, Optional, TextIO, Tuple, Union 11 | 12 | tees = [] 13 | 14 | 15 | @dataclass 16 | class Tee: 17 | orig_stdout: TextIO = field(init=False) 18 | stdout_pipe_proc: Optional[Tuple[TextIO, subprocess.Popen]] = field(init=False, default=None) 19 | 20 | orig_stderr: TextIO = field(init=False) 21 | stderr_pipe_proc: Optional[Tuple[TextIO, subprocess.Popen]] = field(init=False, default=None) 22 | 23 | def __post_init__(self): 24 | tees.append(self) # Do not let this get GC'd 25 | self.orig_stdout = _dup(sys.stdout) 26 | self.orig_stderr = _dup(sys.stderr) 27 | 28 | def to(self, stdout, stderr): 29 | if not isinstance(stdout, list): 30 | stdout = [stdout] 31 | if not isinstance(stderr, list): 32 | stderr = [stderr] 33 | 34 | old_stdout_pipe_proc = self.stdout_pipe_proc 35 | self.stdout_pipe_proc = _tee(self.orig_stdout, stdout, stdout=self.orig_stdout) 36 | self.stdout = stdout 37 | old_stderr_pipe_proc = self.stderr_pipe_proc 38 | self.stderr_pipe_proc = _tee(self.orig_stderr, stderr, stdout=self.orig_stdout) 39 | self.stderr = stderr 40 | 41 | self._drain(stdout_pipe_proc=old_stdout_pipe_proc, stderr_pipe_proc=old_stderr_pipe_proc) 42 | self.resume() 43 | 44 | return self 45 | 46 | def close(self): 47 | self.pause() 48 | self._drain(self.stdout_pipe_proc, self.stderr_pipe_proc) 49 | 50 | def _drain(self, stdout_pipe_proc, stderr_pipe_proc): 51 | # One sharp edge is that if you've spawned a subprocess with 52 | # the redirected stdout/stderr, the tee processes will not 53 | # die. In that case maybe we should set a timeout, or just 54 | # leak them? Not sure. 55 | if stdout_pipe_proc is not None: 56 | pipe, proc = stdout_pipe_proc 57 | pipe.close() 58 | try: 59 | # TODO: replace tee with something that is guaranteed 60 | # to flush on exit 61 | os.kill(proc.pid, signal.SIGINT) 62 | except ProcessLookupError: 63 | pass 64 | proc.wait() 65 | if stderr_pipe_proc is not None: 66 | pipe, proc = stderr_pipe_proc 67 | pipe.close() 68 | try: 69 | os.kill(proc.pid, signal.SIGINT) 70 | except ProcessLookupError: 71 | pass 72 | proc.wait() 73 | 74 | def resume(self): 75 | os.dup2(self.stdout_pipe_proc[0].fileno(), sys.stdout.fileno()) 76 | os.dup2(self.stderr_pipe_proc[0].fileno(), sys.stderr.fileno()) 77 | 78 | def pause(self): 79 | os.dup2(self.orig_stdout.fileno(), sys.stdout.fileno()) 80 | os.dup2(self.orig_stderr.fileno(), sys.stderr.fileno()) 81 | 82 | 83 | def _tee(src, to, stdout): 84 | if src.isatty(): 85 | w, r = os.openpty() 86 | flags = termios.tcgetattr(src.fileno()) 87 | termios.tcsetattr(w, termios.TCSANOW, flags) 88 | termios.tcsetattr(r, termios.TCSANOW, flags) 89 | else: 90 | r, w = os.pipe() 91 | 92 | r = os.fdopen(r, "rb", buffering=0) 93 | w = os.fdopen(w, "wb", buffering=0) 94 | 95 | if src.isatty(): 96 | tty.setraw(r) 97 | 98 | if sys.stdin.isatty(): 99 | # Copy window size 100 | packed = fcntl.ioctl(sys.stdin, termios.TIOCGWINSZ, struct.pack("HHHH", 0, 0, 0, 0)) 101 | fcntl.ioctl(r, termios.TIOCSWINSZ, packed) 102 | 103 | def set_ctty(): 104 | signal.signal(signal.SIGINT, signal.SIG_IGN) 105 | w.close() 106 | fcntl.ioctl(r, termios.TIOCSCTTY, 0) 107 | 108 | else: 109 | 110 | def set_ctty(): 111 | signal.signal(signal.SIGINT, signal.SIG_IGN) 112 | w.close() 113 | 114 | for path in to: 115 | os.makedirs(os.path.dirname(path), exist_ok=True) 116 | 117 | # TODO: fast exit 118 | proc = subprocess.Popen( 119 | ["parent-lifetime", "--term", "tee", "-a"] + list(to), 120 | stdin=r, 121 | start_new_session=True, 122 | stderr=subprocess.DEVNULL, 123 | stdout=stdout, 124 | preexec_fn=set_ctty, 125 | ) 126 | r.close() 127 | return w, proc 128 | 129 | 130 | def _dup(src): 131 | src = os.dup(src.fileno()) 132 | src = os.fdopen(src, "rb", buffering=0) 133 | return src 134 | 135 | 136 | def tee(stdout: Union[str, List[str]], stderr: Union[List[str], str]): 137 | return Tee().to(stdout=stdout, stderr=stderr) 138 | 139 | 140 | if __name__ == "__main__": 141 | t = Tee() 142 | t.to("/tmp/out.orig", "/tmp/err.orig") 143 | t.to("/tmp/out.new", "/tmp/err.new") 144 | breakpoint() 145 | --------------------------------------------------------------------------------