python - 替换 Python 中的所有函数调用

Question

我正在处理一段如下所示的代码:

# This code is not modifiable

from package import distance as dist

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

我的代码:

from package import A

a = A()
a.calculate()

如您所见，distance() 函数被导入到代码的顶部。 A 类调用了一个distance() 函数。它在多个地方这样做，而不仅仅是在 calculate() 中。

我希望类(class)使用我的自定义距离函数。但是，该类不允许我在构造函数中传递它，我无法修改 A 的代码。我该怎么做？这可以通过子类化吗？我尝试了以下方法，但没有用:

from package import A

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

b = B()
b.calculate()

Answer 1

您可以使用 mock.patch功能如下:

距离.py:

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

我的距离.py:

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

a.py:

from distance import distance as dist


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

主要.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()

输出是:

mydistance called

根据您的用例，您可能希望将 with 语句放在其他地方(例如调用站点)。例如:

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

打补丁会产生一些开销。另一种解决方案(基本上与 oetoni 建议的相同)是重新分配属性(记住 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