0

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.

pbreach
  • 16,049
  • 27
  • 82
  • 120
  • 1
    *Why* do you want to do this? What's the actual problem you're trying to solve? It seems like a recipe for unreadable code to me. – jonrsharpe Sep 20 '16 at 17:20
  • 1
    Closely related: [Getting the block of commands that are to be executed in the with statement](http://stackoverflow.com/q/12485837); the solutions there would apply to a decorator as well, as `arg3` is only going to be found as a global. – Martijn Pieters Sep 20 '16 at 17:33
  • I'm trying to allow users to define a system of ordinary differential equations as a series of functions to be integrated in a way that is similar to existing software. – pbreach Sep 20 '16 at 17:36
  • The decorator would bring in some slice objects or constants to be used with numpy arrays which specify a particular element or dimension to subset the array rather than hard coding the slice or element if that makes sense. But you're right, there's probably a better way to do this. – pbreach Sep 20 '16 at 17:37

2 Answers2

1

You can do it like that

from functools import wraps

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

@addarg(3)
def func(arg1, arg2, *args):
    return arg1 + arg2 + sum(args)


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

So that you are not attached to any name and you can apply as many decorators as you want. So if you want to add more decorators it would work. For example

@addarg(3)
@addarg(3)
@addarg(3)
def func(arg1, arg2, *args):
    return arg1 + arg2 + sum(args)

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

# Output would be 12(1 + 2 + 3 + 3 + 3)
vishes_shell
  • 22,409
  • 6
  • 71
  • 81
0

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
martineau
  • 119,623
  • 25
  • 170
  • 301
  • This works for me! It does seem really hacky to use `exec` but the question is a little bit hacky as well. Will wait to see if any other answers come up. – pbreach Sep 21 '16 at 02:01
  • There are even more hacky ways to do it, such as modifying the decorated function's byte-code. There's an example of doing something like that in an [answer](http://stackoverflow.com/a/3454053/355230) to a question about having a `self` argument automatically added to methods of a class (although not necessarily by using a decorator). One advantage of that approach is the result would run faster since it so low-level. – martineau Sep 21 '16 at 02:47