从装饰器中访问自身

Question

在unittest的setUp()方法中，我设置了一些self变量，这些变量在实际的测试中会被引用。我还创建了一个装饰器来进行一些日志记录。有没有办法让我在装饰器中访问这些self变量呢？

为了简单起见，我在这里贴出这段代码：

def decorator(func):
    def _decorator(*args, **kwargs):
        # access a from TestSample
        func(*args, **kwargs)
    return _decorator

class TestSample(unittest.TestCase):    
    def setUp(self):
        self.a = 10

    def tearDown(self):
        # tear down code

    @decorator
    def test_a(self):
        # testing code goes here

从装饰器中访问在setUp()中设置的a，最好的方法是什么呢？

Answer 1

如果你在使用基于类的装饰器实现，这里有一个例子。

import functools
from typing import Any, Callable


class UpdateBackupStatus:

    def __call__(self, f: Callable) -> Any:
        @functools.wraps(wrapped=f)
        def wrapper(calling_instance, *args, **kwargs):
            batch_id = calling_instance.batch_id
            if batch_id is None:
                raise ValueError(
                    "batch_id is required"
                )
            # pre-logic goes here
            response = f(calling_instance, *args, **kwargs)
            # post-logic goes here
            return response

我个人不太喜欢在这种情况下用 self 作为变量名来指代原始调用者，因为这样可能会影响可读性。相反，建议用一个不同且独特的名字来命名（在这个例子中，我用了 calling_instance）。

使用方法如下：

class SomeOtherClass:
    batch_id: str

    def __init__(
        self,
        batch_id: str,
    ):
        self.batch_id = batch_id

    @UpdateBackupStatus()
    def backup_table(self, table_name: str):
        # do something
        pass

现在可以从装饰器中访问 SomeOtherClass 的 batch_id 属性了。

Answer 2

其实你还可以使用 functools.wraps 来保存 help(obj.method) 的信息。除此之外，如果这个装饰器只在类里面使用的话，可以直接放在类的主体里：

from functools import wraps


class MyClass:

    def decorator(func):
        @wraps(func)
        def _decorator(self, *args, **kwargs):
            print("In decorator. self: ", self)
            return func(self, *args, **kwargs)

        return _decorator

    @decorator
    def somemethod(self, a: int, b: int = 5):
        """This is somemethod.

        Parameters
        ----------
        a:
            This is a.
        b:
            The b is optional argument. (Default: 5)
        """
        print("in somemethod")

这样做的效果是这样的：

>>> obj = MyClass()
>>> obj.somemethod()
In decorator. self:  <__main__.MyClass object at 0x7f0378df6920>
in somemethod

调用 help() 时会打印出原来的文档字符串：

>>> help(obj.somemethod)

Help on method somemethod in module __main__:

somemethod(a: int, b: int = 5) method of __main__.MyClass instance
    This is somemethod.
    
    Parameters
    ----------
    a:
        This is a.
    b:
        The b is optional argument. (Default: 5)
:

而如果不加 @wraps(func)，就会变成这样：

>>> help(obj.somemethod)

Help on method _decorator in module __main__:

_decorator(*args, **kwargs) method of __main__.MyClass instance
(END)

Answer 3

因为你是在给一个方法加装饰器，而self是这个方法的一个参数，所以你的装饰器在运行时可以访问到self。显然，在解析时是不能访问的，因为那时候还没有对象，只有一个类。

所以你需要把你的装饰器改成：

def decorator(func):
    def _decorator(self, *args, **kwargs):
        # access a from TestSample
        print 'self is %s' % self
        return func(self, *args, **kwargs)
    return _decorator