tywtyw2002
tywtyw2002

Reputation: 167

How to use decorator to bind a argument to a staticmethod function?

My problem is look like below:

class foo(obj):

    def __init__(self, st='123'):
       self.st = st

    def process(self, x):
       self.st += x

    @staticmethod
    def do_foo(x, myfoo=None):
        if myfoo is None:
           myfoo = foo()
        myfoo.process(x)

def wrapper(fn, st):

    foo_func = foo(st)
    foo.do_foo = functools.partial(foo.do_foo, myfoo=foo_func)
    fn()
    print foo_func.st

    return wrap

@wrapper('stst')
def pro():
    foo.do_foo('double')

def pro2():
    foo.do_foo('double')

pro2()   # <--- normal foo.do_foo
pro()    # <--- partialed foo.do_foo
pro2()   # <--- partialed foo.do_foo

I want to create wrapper decorator to wrapper the staticmethod foo.do_foo with customize foo class, and after the pro() executed, this decorator be able to across the foo object to do some job. ie save the variable value.

How every, in upper code, the wrapper change the foo.do_foo forever globally, not just change it in the decorator scope.

So how to just let foo.do_foo only changed in the decorator scope not global?

Upvotes: 1

Views: 389

Answers (1)

martineau
martineau

Reputation: 123403

Here's an slightly different answer based on the gist.github.com code linked to in your comment (which appears to be based on the first version of my answer).

As I said initially, it sounds to me like what you need to do is make wrapper() a decorator factory function rather than a decorator itself — in other words make it a function that creates and returns a decorator based on its argument(s).

As mentioned in my response to your comment, the problem is that Prof.do staticmethod is essentially a global variable, and if the wrapped function changes it, this will affect all subsequent calls to it — which is the root cause of your problem.

A workaround is to make the wrapped() function the decorator creates save the value of Prof.do before calling the decorated function, and then restore it afterward, so the changes only affect calls to the function made though it. This prevents what it does to Prof.do from messing up other calls to it in other wrapped functions that may be created. It also prevents the accumulation of their effects.

I've encapsulated the changing and restoring the staticmethod by putting it in a contextmanager helper function. One downside of needing to do this is that it increases the amount of overhead involved to make calls to the wrapped function.

import contextlib
import functools

class Prof(object):
    def __init__(self, p='s'):
        self.p = p

    @staticmethod
    def do(x, obj=None):
        if obj is None:
            obj = Prof()
        obj.dprint(x)
        print

    def dprint(self, x):
        print self.p, x
        self.p += x

def wrapper(st):
    @contextlib.contextmanager
    def prof_context(obj):  # could also be defined outside of wrapper function
        # save current staticmethod and replace it with partial below
        saved_method, Prof.do = Prof.do, functools.partial(Prof.do, obj=obj)
        yield
        # undo staticmethod modification
        Prof.do = staticmethod(saved_method)

    def decorator(fn):
        @functools.wraps(fn)
        def wrapped():
            obj = Prof(st)
            print 'current: obj.p is %r' % obj.p
            with prof_context(obj):
                fn()

        return wrapped

    return decorator

def do_p():
    Prof.do('do')

@wrapper('do_p2')
def do_p2():
    Prof.do('do2')

print '#A do_p():'
do_p()
print '#B do_p2():'
do_p2()
print '#C do_p():'
do_p()
print '#D do_p2():'
do_p2()
print '#E do_p():'
do_p()
print '#F do_p():'
do_p()

Output:

#A do_p():
s do

#B do_p2():
current: obj.p is 'do_p2'
do_p2 do2

#C do_p():
s do

#D do_p2():
current: obj.p is 'do_p2'
do_p2 do2

#E do_p():
s do

#F do_p():
s do

Upvotes: 1

Related Questions