计算装饰器内部可能调用或不调用的方法的调用次数

Count calls of a method that may or may not be called inside a decorator

我有这个class:

class SomeClass(object):
    def __init__(self):
        self.cache = {}
    def check_cache(method):
        def wrapper(self):
            if method.__name__ in self.cache:
                print('Got it from the cache!')
                return self.cache[method.__name__]
            print('Got it from the api!')
            self.cache[method.__name__] = method(self)
            return self.cache[method.__name__]
        return wrapper
    @check_cache
    def expensive_operation(self):
        return get_data_from_api()

def get_data_from_api():
    "This would call the api."
    return 'lots of data'

我的想法是,我可以使用 @check_cache 装饰器来防止 expensive_operation 方法在结果已被缓存的情况下额外调用 api 次。

这似乎很好用。

>>> sc.expensive_operation()
Got it from the api!
'lots of data'
>>> sc.expensive_operation()
Got it from the cache!
'lots of data'

但我希望能够使用另一个装饰器对其进行测试:

import unittest

class SomeClassTester(SomeClass):
    def counted(f):
        def wrapped(self, *args, **kwargs):
            wrapped.calls += 1
            return f(self, *args, **kwargs)
        wrapped.calls = 0
        return wrapped
    @counted
    def expensive_operation(self):
        return super().expensive_operation()


class TestSomeClass(unittest.TestCase):
    def test_api_is_only_called_once(self):
        sc = SomeClassTester()
        sc.expensive_operation()
        self.assertEqual(sc.expensive_operation.calls, 1) # is 1
        sc.expensive_operation()
        self.assertEqual(sc.expensive_operation.calls, 1) # but this goes to 2

unittest.main()

问题是 counted 装饰器计算 wrapper 函数被调用的次数,而不是这个内部函数。

如何从 SomeClassTester 算起?

没有简单的方法可以做到这一点。您当前的测试以错误的顺序应用装饰器。你想要 check_cache(counted(expensive_operation)),但你在外面得到了 counted 装饰器:counted(check_cache(expensive_operation)).

counted 装饰器中没有简单的方法来解决这个问题,因为在它被调用时,原始函数已经被 check_cache 装饰器包装了,而且没有简单的方法更改包装器(它在闭包单元中保存对原始函数的引用,从外部是只读的)。

使其工作的一种可能方法是按所需顺序使用装饰器重建整个方法。您可以从闭包单元中获取对原始方法的引用:

class SomeClassTester(SomeClass):
    def counted(f):
        def wrapped(self, *args, **kwargs):
            wrapped.calls += 1
            return f(self, *args, **kwargs)
        wrapped.calls = 0
        return wrapped
    expensive_operation = SomeClass.check_cache(
        counted(SomeClass.expensive_operation.__closure__[0].cell_value)
    )

这当然远非理想,因为您需要确切地知道在 SomeClass 中的方法上应用了哪些装饰器,以便再次正确应用它们。您还需要了解这些装饰器的内部结构,以便您可以获得正确的闭包单元(如果其他装饰器更改为不同,[0] 索引可能不正确)。

另一种(也许更好)的方法可能是更改 SomeClass,这样您就可以在更改的方法和您想要计算的昂贵位之间注入您的计数代码。例如,您可以将真正昂贵的部分放在 _expensive_method_implementation 中,而装饰的 expensive_method 只是调用它的简单包装器。测试 class 可以用它自己的修饰版本覆盖 _implementation 方法(它甚至可能跳过实际昂贵的部分而只是 return 虚拟数据)。它不需要重写常规方法或弄乱它的装饰器。

如果不修改基 class 以提供挂钩或根据基 class 的内部知识更改派生 class 中的整个装饰函数,这是不可能做到的。虽然有第三种基于缓存装饰器内部工作的方法,但基本上改变你的缓存字典以便它计数

class CounterDict(dict):
  def __init__(self, *args):
    super().__init__(*args)
    self.count = {}

  def __setitem__(self, key, value):
    try:
      self.count[key] += 1
    except KeyError:
      self.count[key] = 1
    return super().__setitem__(key, value)


class SomeClassTester(SomeClass):
    def __init__(self):
      self.cache = CounterDict()

class TestSomeClass(unittest.TestCase):
    def test_api_is_only_called_once(self):
        sc = SomeClassTester()
        sc.expensive_operation()
        self.assertEqual(sc.cache.count['expensive_operation'], 1) # is 1
        sc.expensive_operation()
        self.assertEqual(sc.cache.count['expensive_operation'], 1) # is 1