CODEWITHSUNDEEP
pythonargumentsdecorator
Macad_hack
Macad_hack

Reputation: 571

Decorating a function to add custom arguments to function

I am interested in creating a generic function and decorating all other functions with it. The new function will contain arguments from both functions. Example

def decorator(func):
    # Some chunk of code

@decorator
def func(a, b):
    print a**2, b**2
    return a**2, b**2

# Equivalent func of what I want out of the decorator
def equiv_func(a, b, var):
    print var  # I want to modify all decorated functions to accept a new
               # argument and print it using "print var"
    print a**2, b**2
    return a**2, b**2

This is what I kind of came up with but...

def decor(func):
    def modify(var, **kwargs):
       print var
       x, y = func(**kwargs)
       return x, y
    return modify

@decor
def func(a, b):
    print a**2, b**2
    return a**2, b**2

>>> func(var="hi", a=4, b=5)
hi
16 25

I used **kwargs as a way to account for arguments of func, but this forces me to use func(var="hi", a=4, b=5) instead of func(a=4, b=5, var="hi") due to the positioning of **kwargs.

Also, this workaround prevents me from using **kwargs when defining my decorator when I need it since it is already used to input arguments into func. Example

def test(**kwargs):
    if 'test_ans' in kwargs:
        if kwargs['test_ans'] == 'y':
            return True
    else:
        return False

def decor(func):
    def modify(var, **kwargs):
       if test(**kwargs):  # I need to use **kwargs here, but I also need it to define my  
                           # arguments in func.
           print var
       x, y = func(**kwargs)
       return x, y
    return modify

...
# The following is what I expect but due to **kwargs being used for two purposes, 
# it doesn't work
>>> func(var='hi', a=4, b=5, test_ans='y')
'hi'
16 25
>>> func(var='hi', a=4, b=5, test_ans='n')
16 25

The examples are just to illustrate of the constraints I face. They are not the actual code (I wouldn't write a test function that uses **kwargs for example).

To put simply, is there any other way to modify func(a, b) to func(a, b, c, d, e, ..., **kwargs) without using kwargs in the decorator?

Upvotes: 2

Views: 7685

Answers (2)

Patch Rick Walsh
Patch Rick Walsh

Reputation: 457

You can access the parameters that a function accepts with func.func_code.co_varnames. Once you know that, you can filter everything else out.

This is limited in that you would not be able to use this effectively if myfunc accepted **kwargs as a parameter.

def print_decorator(func):
    def wrapped(*args, **kwargs):
        # a tuple of the names of the parameters that func accepts
        func_params = func.func_code.co_varnames
        # grab all of the kwargs that are not accepted by func
        extra = set(kwargs.keys()) - set(func_params)

        for kw in extra:
            print kwargs.pop(kw)

        return func(*args, **kwargs)

    return wrapped

@print_decorator
def myfunc(a, b, c=None):
    print 'myfunc was passed: a={}, b={}, c={}'.format(a, b, c)
    return a, b, c

>>> myfunc(1, b=2, c=3, d=4, e=5)
4
5
myfunc was passed: a=1, b=2, c=3

Upvotes: 4

unutbu
unutbu

Reputation: 879143

You could allow func to be called with

func(a=4, b=5, var="hi")

by making var a keyword argument. 

def decor(func):
    def modify(*args, **kwargs):
        var = kwargs.pop('var', None)
        print var
        x,y=func(*args, **kwargs)
        return x,y
    return modify

@decor
def func(a,b):
    print a**2,b**2
    return a**2,b**2

func(a=4, b=5, var="hi")
func(a=4, b=5)

Upvotes: 2

Related Questions