Using Loops to decorate multiple imported functions in Python

Question

I am new to Python and decorators, so apologies if this seems to be a trivial question.

I am trying to apply decorators to multiple imported functions using a loop in Python as shown below

from random import random, randint, choice

def our_decorator(func):
    def function_wrapper(*args, **kwargs):
        print("Before calling " + func.__name__)
        res = func(*args, **kwargs)
        print(res)
        print("After calling " + func.__name__)
    return function_wrapper

for f in [random, randint, choice]:
    f = our_decorator(f)

random()
randint(3, 8)
choice([4, 5, 6])

Ideally, I expect the output to be in this form :

Before calling random
<random_value>
After calling random
Before calling randint
<random_integer>
After calling randint
Before calling choice
<random_choice>
After calling choice

But, I am only getting the result of the choice function as the output.

<random_choice among 4,5 6>

The decorator has not been applied to any of the functions and it also looks like the random() and the randint(3,8) calls are not getting executed.

I would like to know, what is going wrong here and what can be done to decorate multiple imported functions using loops?

Thanks for the help

Answer 1

Your are decorating the functions and then binding the name f to each of them in turn. So after the loop, f will be equal to the last decorated function. The original names for the original functions are unaffected.

You'll have to work with the names, something like

gl = globals()
for f_name in ['random', 'randint', 'choice']:
    gl[f_name] = our_decorator(gl[f_name])

But I would much prefer to just apply the decorator to each function in turn by hand.

Answer 2

I agree with Remco & Willem that using the @ syntax in the usual way would be better, although Willem's approach of modifying the attributes of the imported random module is probably better than man-handling globals(). But here's another way:

from random import random, randint, choice

def our_decorator(func):
    def function_wrapper(*args, **kwargs):
        print("Before calling " + func.__name__)
        res = func(*args, **kwargs)
        print(res)
        print("After calling " + func.__name__)
    return function_wrapper

random, randint, choice = [our_decorator(f) for f in (random, randint, choice)]

random()
randint(3, 8)
choice([4, 5, 6])

output

Before calling random
0.8171920550436872
After calling random
Before calling randint
8
After calling randint
Before calling choice
4
After calling choice

---

In the comments, RemcoGerlich points out that Willem Van Onsem's technique of decorating the attributes of the random module means that any other imported modules that use the decorated functions in random will also be affected. That's perfectly correct if the other module uses the standard import random statement. I guess that in some circumstances that behaviour may actually be desirable.

However, there is a way that if you have control over the code in the other module. If the other module uses the from random import choice form, then it will get the undecorated version of choice. Here's a short demo.

dectest.py

#!/usr/bin/env python3

from random import choice

def test(seq):
    print('In test')
    return choice(seq)

import random
import dectest

def our_decorator(func):
    def function_wrapper(*args, **kwargs):
        print("Before calling " + func.__name__)
        res = func(*args, **kwargs)
        print(res)
        print("After calling " + func.__name__)
        return res
    return function_wrapper

f = 'choice'
setattr(random, f, our_decorator(getattr(random, f)))

a = [4, 5, 6, 7]
print(random.choice(a))

print('dectest')
print(dectest.test(a))

typical output

Before calling choice
6
After calling choice
6
dectest
In test
7

---

If dectest.py looks like this:

import random

def test(seq):
    print('In test')
    return random.choice(seq)

then we get the behaviour that Remco mentions:

Before calling choice
5
After calling choice
5
dectest
In test
Before calling choice
4
After calling choice
4

Answer 3

You can do this by setting it for the random package like:

import random

for f in ['random','randint','choice']:
    setattr(random,f,our_decorator(getattr(random,f)))

Here you thus set the "attribute" of the random package. Also note that in your for loop, you feed strings.

and then call with:

random.random()
random.randint(3, 8)
random.choice([4, 5, 6])

Nevertheless, this does not look very elegant. A decorator is usually applied by using the @-syntax on the function itself.