从装饰器获取Python函数所属的类

Question

我在Python里有一个装饰器。这个装饰器是一个方法，它接受一个函数作为参数。我想根据传入的函数来创建一个目录结构。我用模块名作为父目录，但想用类名作为子目录。不过，我不知道怎么才能获取到拥有这个函数对象的类的名字。

我的装饰器：

def specialTest(fn):
    filename = fn.__name__
    directory = fn.__module__
    subdirectory = fn.__class__.__name__ #WHERE DO I GET THIS

Answer 1

获取类名

如果你只想要类名（而不是类本身），你可以通过函数的一个属性来获取，这个属性叫做 __qualname__。

import os.path

def decorator(fn):
    filename = fn.__name__
    directory = fn.__module__
    subdirectory = fn.__qualname__.removesuffix('.' + fn.__name__).replace('.', os.path.sep)
    return fn

class A(object):
    @decorator
    def method(self):
        pass
    
    class B(object):
        @decorator
        def method(self):
            pass

如果这个方法所在的类是一个内部类，那么类名会包含外部类的名字。上面的代码通过把所有的点（.）替换成本地路径分隔符来处理这个问题。

当装饰器被调用时，类的其他信息是无法访问的，因为类本身还没有定义。

在方法调用时获取类

如果你需要类本身，并且可以等到装饰的方法第一次被调用时再获取，装饰器可以像往常一样包装这个函数，然后包装器就可以访问实例和类。如果这个方法只需要调用一次，包装器还可以在调用后自己移除，解除装饰。

import types

def once(fn):
    def wrapper(self, *args, **kwargs):
        # do something with the class
        subdirectory = type(self).__name__
        ...
        # undecorate the method (i.e. remove the wrapper)
        setattr(self, fn.__name__, types.MethodType(fn, self))
        
        # invoke the method
        return fn(self, *args, **kwargs)
    return wrapper

class A(object):
    @once
    def method(self):
        pass

a = A()
a.method()
a.method()

注意，这种方法只有在方法被调用时才有效。

在类定义后获取类

如果你需要获取类的信息，即使装饰的方法没有被调用，你可以在包装器上保存一个对装饰器的引用，然后在所有类的方法中查找那些引用了这个装饰器的方法（在相关类定义之后）：

def decorator(fn):
    def wrapper(self, *args, **kwargs):
        return fn(self, *args, **kwargs)
    wrapper.__decorator__ = decorator
    wrapper.__name__ = 'decorator + ' + fn.__name__
    wrapper.__qualname__ = 'decorator + ' + fn.__qualname__
    return wrapper

def methodsDecoratedBy(cls, decorator):
    for method in cls.__dict__.values():
        if     hasattr(method, '__decorator__') \
           and method.__decorator__ == decorator:
            yield method

#...
import sys, inspect

def allMethodsDecoratedBy(decorator)
    for name, cls in inspect.getmembers(sys.modules, lambda x: inspect.isclass(x)):
        for method in methodsDecoratedBy(cls, decorator):
            yield method

这基本上让装饰器在编程中成为一种注解（而不是指Python中的函数注解，它们只用于函数参数和返回值）。一个问题是，装饰器必须是最后一个应用的，否则类属性不会存储相关的包装器，而是另一个外部的包装器。你可以通过在包装器上存储（并稍后检查）所有装饰器来部分解决这个问题：

def decorator(fn):
    def wrapper(self, *args, **kwargs):
        return fn(self, *args, **kwargs)
    wrapper.__decorator__ = decorator
    if not hasattr(fn, '__decorators__'):
        if hasattr(fn, '__decorator__'):
            fn.__decorators__ = [fn.__decorator__]
        else:
            fn.__decorators__ = []
    wrapper.__decorators__ = [decorator] + fn.__decorators__
    wrapper.__name__ = 'decorator(' + fn.__name__ + ')'
    wrapper.__qualname__ = 'decorator(' + fn.__qualname__ + ')'
    return wrapper

def methodsDecoratedBy(cls, decorator):
    for method in cls.__dict__.values():
        if hasattr(method, '__decorators__') and decorator in method.__decorators__:
            yield method

此外，任何你无法控制的装饰器也可以通过装饰它们来合作，这样它们就会在自己的包装器上存储自己，就像 decorator 所做的那样：

def bind(*values, **kwvalues):
    def wrap(fn):
        def wrapper(self, *args, **kwargs):
            nonlocal kwvalues
            kwvalues = kwvalues.copy()
            kwvalues.update(kwargs)
            return fn(self, *values, *args, **kwvalues)
        wrapper.__qualname__ = 'bind.wrapper'
        return wrapper
    wrap.__qualname__ = 'bind.wrap'
    return wrap

def registering_decorator(decorator):
    def wrap(fn):
        decorated = decorator(fn)
        decorated.__decorator__ = decorator
        if not hasattr(fn, '__decorators__'):
            if hasattr(fn, '__decorator__'):
                fn.__decorators__ = [fn.__decorator__]
            else:
                fn.__decorators__ = []
        if not hasattr(decorated, '__decorators__'):
            decorated.__decorators__ = fn.__decorators__.copy()
        decorated.__decorators__.insert(0, decorator)
        decorated.__name__ = 'reg_' + decorator.__name__ + '(' + fn.__name__ + ')'
        decorated.__qualname__ = decorator.__qualname__ + '(' + fn.__qualname__ + ')'
        return decorated
    wrap.__qualname__ = 'registering_decorator.wrap'
    return wrap

class A(object):
    @decorator
    def decorated(self):
        pass
    
    @bind(1)
    def add(self, a, b):
        return a + b
    
    @registering_decorator(bind(1))
    @decorator
    def args(self, *args):
        return args
    
    @decorator
    @registering_decorator(bind(a=1))
    def kwargs(self, **kwargs):
        return kwargs

A.args.__decorators__
A.kwargs.__decorators__
assert not hasattr(A.add, '__decorators__')
a = A()
a.add(2)
# 3

另一个问题是扫描所有类的效率不高。你可以通过使用一个额外的类装饰器来注册所有类，以检查方法装饰器，从而提高效率。然而，这种方法比较脆弱；如果你忘记给类加装饰器，它就不会被记录在注册表中。

class ClassRegistry(object):
    def __init__(self):
        self.registry = {}
    
    def __call__(self, cls):
        self.registry[cls] = cls
        cls.__decorator__ = self
        return cls
    
    def getRegisteredClasses(self):
        return self.registry.values()

class DecoratedClassRegistry(ClassRegistry):
    def __init__(self, decorator):
        self.decorator = decorator
        super().__init__()
    
    def isDecorated(self, method):
        return (    hasattr(method, '__decorators__') \
                and self.decorator in method.__decorators__) \
            or (    hasattr(method, '__decorator__') \
                and method.__decorator__ == self.decorator)
    
    def getDecoratedMethodsOf(self, cls):
        if cls in self.registry:
            for method in cls.__dict__.values():
                if self.isDecorated(method):
                    yield method
    
    def getAllDecoratedMethods(self):
        for cls in self.getRegisteredClasses():
            for method in self.getDecoratedMethodsOf(cls):
                yield method

用法：

decoratedRegistry = DecoratedClassRegistry(decorator)

@decoratedRegistry
class A(object):
    @decoratedRegistry
    class B(object):
        @decorator
        def decorated(self):
            pass
        
        def func(self):
            pass
    
    @decorator
    def decorated(self):
        pass
    
    @bind(1)
    def add(self, a, b):
        return a + b
    
    @registering_decorator(bind(1))
    @decorator
    def args(self, *args):
        return args
    
    @decorator
    @registering_decorator(bind(a=1))
    def kwargs(self, **kwargs):
        return kwargs

decoratedRegistry.getRegisteredClasses()
list(decoratedRegistry.getDecoratedMethodsOf(A.B))
list(decoratedRegistry.getDecoratedMethodsOf(A))
list(decoratedRegistry.getAllDecoratedMethods())

监控多个装饰器和应用多个装饰器注册的部分留作练习。

Answer 2

在Python 2中，你可以在方法对象上使用im_class这个属性。而在Python 3中，你需要用__self__.__class__（或者type(method.__self__)）来代替。

Answer 3

如果 fn 是一个 实例方法，那么你可以使用 fn.im_class 来获取它所属的类。

>>> class Foo(object):
...     def bar(self):
...         pass
...
>>> Foo.bar.im_class
__main__.Foo

需要注意的是，这种方法在装饰器中是 不 可行的，因为一个函数只有在类定义之后才会变成实例方法（也就是说，如果你用 @specialTest 来装饰 bar，那么就无法使用这个方法；如果真的能做到，那也得通过检查调用栈或者其他一些复杂的方式来实现，这样做可就麻烦了）。