CODEWITHSUNDEEP
pythonpython-decorators
edgarstack
edgarstack

Reputation: 1116

Python: Best way to decorate a function only locally

What is the best way to decorate a function, but only in a local scope. For example we have the following:

def a():
    do stuff
    b()
    do stuff

I want to make a function c() that does the same stuff but decorates b so it adds some extra stuff. I found a way to do that but it changes b() globally:

def c():
    global b
    b = decorator(b)
    a()

Is there a way to do that but without changing the function b globally?

Upvotes: 0

Views: 96

Answers (3)

Prashant Sinha
Prashant Sinha

Reputation: 323

You can patch the function b, and replace it with another function. Use unittest.mock in Python 3, or mock in Python 2 for this:

from unittest.mock import patch

def c():
    with patch('b', decorator(b)):
        a()

# Or, as a decorator
@patch('b', decorator(b))
def c():
    a()

By patching, you're essentially replacing the b instance with your own.

Upvotes: 2

eqzx
eqzx

Reputation: 5599

Since you're rebinding b() so that calling a() will respect the updated b(), you can save the original b(), and restore it afterwards. This may not be the most efficient way, but it worked for me.

import copy
def c():
    global b
    orig_b = copy.deepcopy(b)
    b = decorator(b)
    a()
    b = orig_b

Upvotes: 1

Serge Ballesta
Serge Ballesta

Reputation: 149125

Python allow any object to be a callable, not only functions.

You could use something like:

class A(object):
    def __init__(self, func=b):
        self.func = func
    def __call__(self):
        # do stuff
        self.func()
        # do other stuff

Then with a = A(), a() will simply invoke f but if you writec = A(decorator(b)), c() will invoke the decorated version

Upvotes: 1

Related Questions