Search code examples
pythonfunctionnamespacespython-decorators

Using a decorator to bring variables into decorated function namespace


Say I have a function that looks like this:

def func(arg1, arg2):
    return arg1 + arg2 + arg3

And I want to bring in arg3 using an argument passed to a decorator. Is this possible and if so how would this be done?

I've tried this which was suggested here to be able to pass arguments to decorators:

from functools import wraps

def addarg(arg):
    def decorate(func):
        arg3 = arg
        @wraps(func)
        def wrapped(*args):
            return func(*args)
        return wrapped
    return decorate

@addarg(3)
def func(arg1, arg2):
    return arg1 + arg2 + arg3

if __name__ == "__main__":
    print(func(1, 2))

But I get this error (understandably):

Traceback (most recent call last):
  File "C:/Users/pbreach/Dropbox/FIDS/PhD/Code/anemi3/utils.py", line 19, in <module>
    print(func(1, 2))
  File "C:/Users/pbreach/Dropbox/FIDS/PhD/Code/anemi3/utils.py", line 9, in wrapped
    return func(*args)
  File "C:/Users/pbreach/Dropbox/FIDS/PhD/Code/anemi3/utils.py", line 16, in func
    return arg1 + arg2 + arg3
NameError: name 'arg3' is not defined

For my application it would be much nicer to be able to define the function in this way. There are other ways but then more code will have to be added to the body of each function.


Solution

  • It's a little on the hacky side, but this implmentation of the decorator seems to do what you want. The arg3 variable had to be made to appear as a global because it's referenced as one in the byte-code generated from the definition of func(), so it's not really considered a function argument—but that doesn't appear to matter in this case.

    from functools import wraps
    
    def addarg(arg):
        def decorate(func):
            @wraps(func)
            def wrapped(*args):
                global_vars, local_vars = {'func': func, 'args': args}, {}
                func.__globals__.update(arg3=arg)
                exec('_result = func(*args)', global_vars, local_vars)
                return local_vars['_result']
            return wrapped
        return decorate
    
    @addarg(3)
    def func(arg1, arg2):
        return arg1 + arg2 + arg3
    
    if __name__ == "__main__":
        print(func(1, 2))  # -> 6