Suppose that a function contains a lot of print-statements.

I want to capture all of those print statements in a string, or save them to text file.

What kind of function decorator might do that for us?

log_file = open("log.txt", "w")  

@copy_print_statements(log_file)
def printy_the_printer():  
    print("I print a lot")
   
# should print to both `sys.stdout` and `log_file`
printy_the_printer() 
printy_the_printer()
printy_the_printer()

The following is one failed attempt. Feel free to ignore, or depart from the code below. The real goal is to write code for a decorator. The decorator replaces an old function with a new function. The old functions print a lot to console and the new functions send the print-statements somewhere else.

import io
import sys
import functools 

class MergedStream:
    """
    """
    def __init__(self, lefty, righty):
        """
        `lefty` and `righty` should be file-streams.  
  
        Examples of valid streams might be the values returned by
        the following function calls:

            getattr(sys, 'stdout') 
            io.StringIO()
            open("foo.txt", "w")
        """
        self._lefty  = lefty
        self._righty = righty

    def write(self, *args, **kwargs):
        """
        """
        self._lefty.write(*args, **kwargs)
        self._righty.write(*args, **kwargs)

class CopyPrintStatements:
    def __init__(_callable, file):
        self._callable = _callable 
        self._file = _file  

    def __call__(*args, **kwargs):
        old_stdout = sys.stdout  
        sys.stdout = MergedStream(sys.stdout, self._file)
        try:
            return self._callable(*args, **kwargs) 
        finally:
            sys.stdout = old_stdout

    @classmethod
    def copy_print_statements(cls, file_stream):
        """
             This class method is intended to decorate callables
 
             An example usage is shown below:  

             @copy_print_statements(sys.stderr)
             def foobar():
                  print("this message is printed to both `stdout` and `stderr`")
        """ 
        decorator = cls.make_decorator(file_stream)
        return decorator  

    @classmethod
    def make_decorator(cls, old_callable, file):
        new_callable = cls(old_callable, file)
        new_callable = functools.wraps(old_callable) 
        return new_callable  

CodePudding user response：

I think I have made what you wanted, I basically made a context manager that tempoaraly redirects sys.stdout to itself, then it redirects the messages onto the old sys.stdout and the file. I then made a function decorator which uses the context manager when running the functions:

import sys


class Logger:

    def __init__(self, file):
        self.file = file

    def __enter__(self):
        self.terminal = sys.stdout
        self.f = open(self.file, "a")
        sys.stdout = self

    def __exit__(self, *args):
        sys.stdout = self.terminal
        self.f.close()

    def write(self, message):
        self.terminal.write(message)
        self.f.write(message)

    def flush(self):
        self.terminal.flush()
        self.f.flush()


def copy_print_statements(file):
    def copy_print(fn):
        def wraps(*args, **kwargs):
            with Logger(file):
                fn(*args, **kwargs)
        return wraps
    return copy_print


@copy_print_statements("output.txt")
def printy_the_printer(x):
    print(f"I print a lot {x}")


# should print to both `sys.stdout` and `log_file`
printy_the_printer(1)
printy_the_printer(2)
printy_the_printer(3)

I made the context manager append to the file ("a") rather than overwrite ("w")

CodePudding user response：

My 2 cents on the problem. Simplest as possible.

import sys

def copy_print_statements(file):
    def decorator(func):
        def wrapper(*args, **kwargs):
            stdout_backup = sys.stdout
            stderr_backup = sys.stderr
            with open(file, "a") as f:
                sys.stdout = f
                sys.stderr = f
                func(*args, **kwargs)

            sys.stdout = stdout_backup
            sys.stderr = stderr_backup
        return wrapper
    return decorator

@copy_print_statements("log.txt")
def printy_the_printer():  
    print("I print a lot to stdout")
    print("I print a lot to stderr", file=sys.stderr)

printy_the_printer()
printy_the_printer()
printy_the_printer()

Sadly, but idk any way to close the file on exit, so reopen it on each call.