装饰装饰器：如何为装饰器程序化分配功能

你现有的foo、bar这类类别装饰器逻辑重复，尝试用“装饰装饰器”category封装逻辑时失败，问题出在两处：

1. 类注册逻辑检查的是带下划线前缀的属性（比如_foo），但你写的category给函数设置的是不带下划线的属性（比如foo），二者不匹配；
2. 原category的写法冗余，没必要先定义空函数再装饰，直接做装饰器工厂生成装饰器更高效。

正确实现方式 #

方式一：装饰器工厂（最简洁） #

直接传入类别名称，就能生成对应的装饰器，完全避免重复代码：

def category(cat_name):
    # 生成匹配类注册逻辑的下划线前缀属性名
    attr_name = f"_{cat_name}"
    def decorator(func):
        setattr(func, attr_name, None)
        return func
    return decorator

# 生成foo和bar装饰器
foo = category("foo")
bar = category("bar")

方式二：用装饰器语法生成（偏好装饰器写法时用） #

如果你习惯用@category的语法定义装饰器，可以调整category的实现，让它装饰一个占位函数，提取函数名作为类别名：

def category(cat_func):
    cat_name = cat_func.__name__
    attr_name = f"_{cat_name}"  # 补上下划线前缀，匹配注册逻辑
    def decorator(func):
        setattr(func, attr_name, None)
        return func
    return decorator

# 用@category装饰占位函数生成装饰器
@category
def foo():
    pass

@category
def bar():
    pass

验证效果 #

把上面任意一种实现替换掉原有的foo、bar装饰器定义，运行完整代码，输出结果和之前完全一致：

def class_register_categories(cls):
    categories = 'foo', 'bar'
    
    for category in categories:
        setattr(cls, category, Category())
    
    for methodname in dir(cls):
        method = getattr(cls, methodname)
        for category in categories:
            if hasattr(method, '_'+category):
                getattr(cls, category).append(methodname)
        
    return cls

class Category(list):
    def __repr__(self):
        string = '\n'.join(self)
        return string

# 这里替换成上面的任意一种category实现
foo = category("foo")
bar = category("bar")

@class_register_categories
class MyClass(object):
    @foo
    def my_method(self, arg1, arg2):
        pass
    
    @foo
    @bar
    def my_other_method(self, arg1, arg2):
        pass
    
    def yet_another_method(self):
        pass

myclass = MyClass()
print(myclass.foo)
# 输出：
# my_method
# my_other_method

新增类别更省心 #

现在新增类别只需要一行代码，比如加个baz类别：

baz = category("baz")
# 或者用装饰器语法：
@category
def baz():
    pass

别忘了在class_register_categories的categories元组里加上'baz'就行。

内容的提问来源于stack exchange，提问作者Puff