CODEWITHSUNDEEP
python
r0f1
r0f1

Reputation: 3106

Replace all function calls in Python

I am working with a piece of code that looks like this:

# This code is not modifiable

from package import distance as dist

class A:
    def calculate(self):
        ...
        # call to dist()
        ...

My code:

from package import A

a = A()
a.calculate()

As you can see, the distance() function is imported at the top of the code. The class A makes a call to a distance() function. It does so, in several places and not only in calculate().

I want the class to use my custom distance function. However, the class does not let me pass it in the constructor and I cannot modify the code of A. How would I do this? Is this possible via subclassing? I tryed the following, which did not work:

from package import A

class B(A):
    def __init__(self):
        from mypackage import mydistance as dist
        return super().__init__()

b = B()
b.calculate()

Upvotes: 3

Views: 859

Answers (3)

oetoni
oetoni

Reputation: 3887

you can do this immediately from you code without adding new imports or using mock.patch or any third class B. Immediately from yourcode where you are referencing A

Now you just need to do a small change on you code, don't just do from package import A but import directly import baseA that way we control all the variables in it also the imports ;) and as such we can. Then we can use it's local variable name in order to catch it and change it in runtime :D as such:

Supposed distance() function in import of A

distance.py

def distance(something):
    return something + 1

then I assume the name baseA for the file that contains the class A

baseA.py

# This code is not modifiable

from distance import distance as dist


class A:
    def calculate(self):
        something = dist(1)
        return something

finally youcode file that contains you running code and the place from where we want to modify the dist() that it is imported in baseA

yourcode.py

import baseA


def newDist(something):
    return something + 2


baseA.dist = newDist
a = baseA.A()


something = a.calculate()
print(something)

I successfully managed to change the behavior of dist in baseA as such :) baseA.dist = newDist remember not to put parenthesis because we are passing the function as an object in order to assign its behavior to the baseA.dist that is imported from distance.py file

testing the solution

enter image description here

Upvotes: 1

Smi
Smi

Reputation: 14336

You can use the mock.patch function as follows:

distance.py:

def distance():
    print('distance called')

mydistance.py:

def mydistance():
    print('mydistance called')

a.py:

from distance import distance as dist


class A:
    def calculate(self):
        dist()

main.py:

from unittest import mock

from a import A
from mydistance import mydistance


class B(A):
    def calculate(self):
        with mock.patch('a.dist', wraps=mydistance):
            super().calculate()


if __name__ == '__main__':
    b = B()
    b.calculate()

Output is:

mydistance called

Depending on your use case, you might want to put the with statement somewhere else (such as in the call site). For example:

if __name__ == '__main__':
    b = B()
    with mock.patch('a.dist', wraps=mydistance):
        for _ in range(0, 100):
            b.calculate()

Patching causes some overhead. Another solution (basically the same as oetoni suggests) is to reassign the attribute (remember to import a):

if __name__ == '__main__':
    b = B()
    old_dist = a.dist
    a.dist = mydistance
    for _ in range(0, 100):
        b.calculate()
    a.dist = old_dist

Upvotes: 1

René Pijl
René Pijl

Reputation: 4738

You can do this using the patch@ decorator in unittest.mock.

Upvotes: -1

Related Questions