# !/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
# Taken and modified from original source:
# https://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/
import ctypes
import logging
import os
import sys
from contextlib import contextmanager
from functools import partial
IS_WINDOWS = sys.platform == "win32"
IS_MACOS = sys.platform == "darwin"
logger = logging.getLogger(__name__)
def get_libc():
if IS_WINDOWS or IS_MACOS:
logger.warning(
"NOTE: Redirects are currently not supported in Windows or MacOs."
)
return None
else:
return ctypes.CDLL("libc.so.6")
libc = get_libc()
def _c_std(stream: str):
return ctypes.c_void_p.in_dll(libc, stream)
def _python_std(stream: str):
return {"stdout": sys.stdout, "stderr": sys.stderr}[stream]
_VALID_STD = {"stdout", "stderr"}
@contextmanager
def redirect(std: str, to_file: str):
"""
Redirects ``std`` (one of ``"stdout"`` or ``"stderr"``) to a file
in the path specified by ``to_file``. This method redirects the
underlying std file descriptor (not just pyton's ``sys.stdout|stderr``).
See usage for details.
Directory of ``dst_filename`` is assumed to exist and the destination file
is overwritten if it already exists.
.. note:: Due to buffering cross source writes are not guaranteed to
appear in wall-clock order. For instance in the example below
it is possible for the C-outputs to appear before the python
outputs in the log file.
Usage:
::
# syntactic-sugar for redirect("stdout", "tmp/stdout.log")
with redirect_stdout("/tmp/stdout.log"):
print("python stdouts are redirected")
libc = ctypes.CDLL("libc.so.6")
libc.printf(b"c stdouts are also redirected"
os.system("echo system stdouts are also redirected")
print("stdout restored")
"""
if std not in _VALID_STD:
raise ValueError(
f"unknown standard stream <{std}>, must be one of {_VALID_STD}"
)
c_std = _c_std(std)
python_std = _python_std(std)
std_fd = python_std.fileno()
def _redirect(dst):
libc.fflush(c_std)
python_std.flush()
os.dup2(dst.fileno(), std_fd)
with os.fdopen(os.dup(std_fd)) as orig_std, open(to_file, mode="w+b") as dst:
_redirect(dst)
yield
_redirect(orig_std)
redirect_stdout = partial(redirect, "stdout")
redirect_stderr = partial(redirect, "stderr")