CODEWITHSUNDEEP
pythonooploggingoverridingstdout
Arrajj
Arrajj

Reputation: 187

How to override the print function with wrapper while changing stdout

I am trying to build a wrapper that has the following goals:

  1. To catch the prints for any function to save them for later ( let's return a list for the moment )
  2. To override the prints also in the wrapped functions to add a timestamp for better logging and more information ( I don't want to log.info() because this implies to change the prints in the wrapped function, and if I want to do that in the wrapper later I will lose the precision of the time )

the goal is always to not change anything in the wrapped function, aside than only to wrap them

Here is my code so far:

def my_decorator(func):
    def wrapped_func(*args,**kwargs):
         return func("Here we add a timestamp",*args,**kwargs)
    return wrapped_func

def list_prints_to_logs(list_prints):
    for log in list_prints:
        logger.info(log)

def my_handler(func):
    def wrapper(*args,**kwargs):

        print('this is BEFORE the function')
        old_print = print                        
        print = my_decorator(print)

        ### catching the outputs of the function
        old_stdout = sys.stdout
        sys.stdout = tempstdout = StringIO()
        results = func(*args,**kwargs)  # THE ERROR STARTS IN HERE AFTER ENTERING THIS FUNCTION
        sys.stdout = old_stdout


        print('this is AFTER the function')
        prints_in_function = tempstdout.getvalue().splitlines()
        list_prints_to_logs(prints_in_function)  

        #trying to return the print like it was before
        print = old_print 
        print(*prints_in_function, sep='\n')        

        return results
    return wrapper


@my_handler
def my_func(arg=1,kwarg='yes'):

    print('this is inside the Function')
    print('this is second print in the function')


my_func()

and the error that is returned to me is :

UnboundLocalError: local variable 'print' referenced before assignment

the line of the error has started from the print inside my_func

I am facing difficulty in how to override a function temporarily, and persist this override for the next calls,

Upvotes: 1

Views: 543

Answers (1)

Boseong Choi
Boseong Choi

Reputation: 2596

UnboundLocalError is caused when you had assigned local variable but used it before assingment.
Normaly when you use assignment statement in function, it binds a local variable.
If you want to bind gloabl variable, you should use global statement before assinging.

def my_handler(func):
    def wrapper(*args,**kwargs):
        global print
        ...
        print = my_decorator(print)
        ...
        print = old_print 
        print(*prints_in_function, sep='\n')        

        return results
    return wrapper

However, this decorator doen't affect to outside of module. If you want to affect including other modules, you should rebind builtins.print.

import builtins
def my_handle(func):
    def wrapper(*args,**kwargs):
        ...
        builtins.print = my_decorator(print)
        ...
        builtins.print = old_print
        ...
    return wrapper

Upvotes: 2

Related Questions