Overload function in all subfunctions in scope

user_4690023 0
Viewed 65

Let's say that I want to overload a standard function with a customized version, I can simply write

original_function_name = customized_function

For example I can do:

def custom_print(s):
    print(f'!!{s}!!')

def fun1(n):
    print = custom_print
    print(n)

fun1(1)

>> !!1!!

However this override is only valid inside fun1. If I do

def custom_print(s):
    print(f'!!{s}!!')

def fun1(n):
    print = custom_print
    fun2(n)

def fun2(n):
    print(n)

if __name__ == '__main__':
    fun1(1)

>>1

the custom print function is not passed (clearly) to fun2 which uses the standard print function. Is there a way to override the function not just in the scope where I define it, but also in all called function?

NOTE: This is a minimal example, in the real code there are several nested functions imported from different modules and I'd like to override the function in all of them without modifying them one by one.

NOTE2: I recognize this is bad practice and should not be done as it goes against a number of best-practice coding principles.

python-3.x python
3 Answers

Is there a way to override the function not just in the scope where I define it, but also in all called function?

Since Python does not have dynamic scopes, the only way is to swap the print builtin itself, then restore it afterwards. This is wildly unsafe: in a multithreaded environment, or when coroutines (generators) are involved, as they will see the replacement function for the entirety of that swap.

Furthermore the code you show here will obviously not work: you can't use the replaced print, since it's been replaced. And not precisely matching the proper signature of the function you're replacing... is probably a bad idea. That's not hard to solve though: the signature of print is just *args, end='\n', sep=' ', file=sys.stdout, flush=False, and you can file.write(...) with your content.

Then you can just update the __builtins__ dict (this should be available in all scopes):

# swapper.py
import builtins
import contextlib
import sys

def _my_print(*args, end='\n', sep=' ', file=sys.stdout, flush=False):
     file.write('XXX ')
     file.write(sep.join(map(str, args)))
     file.write(end)
     if flush:
          file.flush()

@contextlib.contextmanager
def swap():
     old_print = print
     builtins.print = _my_print
     try:
          yield
     finally:
          builtins.print = old_print

>>> print("ok")
ok
>>> import swapper
>>> print("ok")
ok
>>> with swapper.swap():
...     print('ok')
... 
XXX ok
>>> print('ok')
ok
user_81821180

You can select the scope in which you monkey patch your function. Python evaluates in LEGB order:

  • Local (which you show)
  • Enclosing (the nested functions you mention)
  • Global (module)
  • Builtin (actually just a special module)

You can do print = monkey_print at any of those levels. Just make sure it follows the same interface that all the other parts of your program expect: print(*args, **kwargs) is usually a safe bet.

Here are some examples:

from sys import stdout
import builtins

def bprint(*args, **kwargs):
    kwargs.get('file', stdout).write('Builtin!: ' + kwargs.get('sep', ' ').join(map(str, args)) + kwargs.get('end', '\n'))

def gprint(*args, **kwargs):
    # Otherwise this will be infinite recursion
    builtins.print('Global! ', *args, **kwargs)

def eprint(*args, **kwargs):
    print('Enclosing! ', *args, **kwargs)

def lprint(*args, **kwargs):
    print('Local! ', *args, **kwargs)

builtins.print = bprint
print = gprint

def decorator(func):
    def wrapper(*args, **kwargs):
        print(*args, **kwargs)
        return func(*args, **kwargs)
    print = eprint
    return wrapper

@decorator
def func(*args, **kwargs):
    print = lprint
    print(*args, **kwargs)

print('Example complete')
func('Print me next!')

The output of this script will be

Builtin!: Global!  Example complete
Builtin!: Global!  Enclosing!  Print me next!
Builtin!: Global!  Local!  Print me next!

Is this a good idea? Probably not when it comes to replacing a ubiquitous function like print. However, knowing how to monkey-patch properly is an important tool when it comes to unit testing, among other things.

user_29887300

Easy fix:

def custom_print(s):
    print(f'!!{s}!!')

def fun1(n):
    print = custom_print
    fun2(n, print)

def fun2(n, print = print):
    print(n)

fun1(1)
user_152006130
Related