1# mypy: allow-untyped-defs 2# !/usr/bin/env python3 3 4# Copyright (c) Facebook, Inc. and its affiliates. 5# All rights reserved. 6# 7# This source code is licensed under the BSD-style license found in the 8# LICENSE file in the root directory of this source tree. 9 10# Taken and modified from original source: 11# https://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/ 12import ctypes 13import logging 14import os 15import sys 16from contextlib import contextmanager 17from functools import partial 18 19 20IS_WINDOWS = sys.platform == "win32" 21IS_MACOS = sys.platform == "darwin" 22 23 24logger = logging.getLogger(__name__) 25 26 27def get_libc(): 28 if IS_WINDOWS or IS_MACOS: 29 logger.warning( 30 "NOTE: Redirects are currently not supported in Windows or MacOs." 31 ) 32 return None 33 else: 34 return ctypes.CDLL("libc.so.6") 35 36 37libc = get_libc() 38 39 40def _c_std(stream: str): 41 return ctypes.c_void_p.in_dll(libc, stream) 42 43 44def _python_std(stream: str): 45 return {"stdout": sys.stdout, "stderr": sys.stderr}[stream] 46 47 48_VALID_STD = {"stdout", "stderr"} 49 50 51@contextmanager 52def redirect(std: str, to_file: str): 53 """ 54 Redirect ``std`` (one of ``"stdout"`` or ``"stderr"``) to a file in the path specified by ``to_file``. 55 56 This method redirects the underlying std file descriptor (not just python's ``sys.stdout|stderr``). 57 See usage for details. 58 59 Directory of ``dst_filename`` is assumed to exist and the destination file 60 is overwritten if it already exists. 61 62 .. note:: Due to buffering cross source writes are not guaranteed to 63 appear in wall-clock order. For instance in the example below 64 it is possible for the C-outputs to appear before the python 65 outputs in the log file. 66 67 Usage: 68 69 :: 70 71 # syntactic-sugar for redirect("stdout", "tmp/stdout.log") 72 with redirect_stdout("/tmp/stdout.log"): 73 print("python stdouts are redirected") 74 libc = ctypes.CDLL("libc.so.6") 75 libc.printf(b"c stdouts are also redirected" 76 os.system("echo system stdouts are also redirected") 77 78 print("stdout restored") 79 80 """ 81 if std not in _VALID_STD: 82 raise ValueError( 83 f"unknown standard stream <{std}>, must be one of {_VALID_STD}" 84 ) 85 86 c_std = _c_std(std) 87 python_std = _python_std(std) 88 std_fd = python_std.fileno() 89 90 def _redirect(dst): 91 libc.fflush(c_std) 92 python_std.flush() 93 os.dup2(dst.fileno(), std_fd) 94 95 with os.fdopen(os.dup(std_fd)) as orig_std, open(to_file, mode="w+b") as dst: 96 _redirect(dst) 97 try: 98 yield 99 finally: 100 _redirect(orig_std) 101 102 103redirect_stdout = partial(redirect, "stdout") 104redirect_stderr = partial(redirect, "stderr") 105