├── .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]} <subcommand to run>\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 | 


--------------------------------------------------------------------------------