Canaryyellow
Canaryyellow

Reputation: 167

How to apply a decorator to subroutines in python

I want to modify the way several of my functions behave, so the use of a decorator comes to mind. For example let's say I have a batch data taking function takeData(s):

def takeData(s):
    takeDataSet_1(s)
    takeDataSet_2(s)
    .
    .
    .

A simple thing that I might want to do is update the parameter dict s, before each of the takeDataSet function calls. So that the effective code would be more like this:

def takeData(s):
    s = updateParams(s)
    takeDataSet_1(s)
    s = updateParams(s)
    takeDataSet_2(s)
    s = updateParams(s)
    .
    .
    .

Is there a way to do this with a decorator so that my code would look more like

@takeDataWithUpdatedParams
def takeData(s):
    takeDataSet_1(s)
    takeDataSet_2(s)

Is there a way to control the depth of recursion such a decorator? So that if takeDataSet_1(s) had subroutines of its own s could be updated between them, as in:

@recursiveUpdateParams
def takeData(s):
    takeDataSet_1(s)
    takeDataSet_2(s)

def takeDataSet_1(s):
    takeData_a(s)
    takeData_b(s)

Gets executed as

def takeData(s):
    s = updateParams(s)
    takeDataSet_1(s)
    s = updateParams(s)
    takeDataSet_2(s)
    s = updateParams(s)

def takeDataSet_1(s):
    s = updateParams(s)
    takeData_a(s)
    s = updateParams(s)
    takeData_b(s)

Upvotes: 0

Views: 131

Answers (2)

Huazuo Gao
Huazuo Gao

Reputation: 1733

Very interesting question. To achieve this, you need to dive deep into the function object (without eval, exec or ast stuff, anyway).

def create_closure(objs):
    creat_cell = lambda x: (lambda: x).__closure__[0]
    return tuple(create_cell(obj) for obj in objs)

def hijack(mapper):
    from types import FunctionType
    def decorator(f):
        globals_ = {k: mapper(v) for k, v in f.__globals__.items()}
        closure_ = f.__closure__
        if closure_:
            closure_ = create_closure(i.cell_contents for i in closure_)
        return (lambda *arg, **kwarg:
                FunctionType(f.__code__, globals_, f.__name__,
                             f.__defaults__, closure_)(*arg, **kwarg))
    return decorator

Test:

x = 'x'
y = 'y'
@hijack(lambda obj: 'hijacked!' if obj is x else obj)
def f():
    return (x, y)
f()

Out:

('hijacked!', 'y')

Finally a solution to the original problem:

x = lambda: 'x()'
y = lambda: 'y()'
mydecorator = lambda f: lambda *arg, **kwarg: f(*arg, **kwarg) + ' decorated!'
targets = {id(x): mydecorator(x)}
def mapper(obj):
    if id(obj) in targets:
        return targets[id(obj)]
    else:
        return obj
@hijack(mapper)
def f():
    return (x(), y())
f()

Out:

('x() decorated!', 'y()')

Upvotes: 1

BrenBarn
BrenBarn

Reputation: 251458

No. Decorators wrap a function, which means they can add their own behavior before and after the function. They can't change what happens "in the middle" of the function, as your example does (let alone what happens in the middle of other functions called by that function, as in your recursion example).

You could create a decorator-like function that accepts your takeData_* functions as arguments and does the updating, so you'd do something like:

def updateAndCall(func, params):
    s = updateParams(params)
    func(s)

def takeData(s):
    updateAndCall(takeData_1, s)
    updateAndCall(takeData_2, s)

However, whether this will be useful depends on the interaction between the various functions. In particular, with this approach each "update" happens to the original s; the updates are not cumulative with the updated s being updated again for the second call, etc.

Upvotes: 1

Related Questions