标记函数为可调用的装饰器

Question

我正在创建一个函数标签系统，目的是根据标签来启用或禁用某些功能：

def do_nothing(*args, **kwargs): pass

class Selector(set):
    def tag(self, tag):
        def decorator(func):
            if tag in self:
                return func
            else:
                return do_nothing
        return decorator

selector = Selector(['a'])

@selector.tag('a')
def foo1():
    print "I am called"

@selector.tag('b')
def foo2():
    print "I am not called"

@selector.tag('a')
@selector.tag('b')
def foo3():
    print "I want to be called, but I won't be"

foo1() #Prints "I am called"
foo2() #Does nothing
foo3() #Does nothing, even though it is tagged with 'a'

我想问的是最后一个函数，foo3。我明白为什么它没有被调用。我在想有没有办法让它在选择器中任何标签存在时被调用。理想情况下，解决方案应该只检查一次标签，而不是每次调用函数时都检查。

顺便提一下：我这样做是为了根据环境变量在unittest单元测试中选择要运行的测试。我实际的实现使用了unittest.skip。

编辑：添加了装饰器的返回。

Answer 1

这样做对你有用吗：

class Selector(set):
    def tag(self, tag_list):
        def decorator(func):
            if set(tag_list) & self:
                return func
            else:
                return do_nothing
        return decorator


@selector.tag(['a','b'])
def foo3():
    print "I want to be called, but I won't be"

Answer 2

问题是，如果你把它装饰两次，一个会返回函数，另一个什么都不返回。

foo3() -> @selector.tag('a') -> foo3()
foo3() -> @selector.tag('b') -> do_nothing

foo3() -> @selector.tag('b') -> do_nothing
do_nothing -> @selector.tag('a') -> do_nothing

这就意味着，不管你怎么顺序调用，你总是得不到任何东西。你需要做的是在每个对象上保持一组标签，并一次性检查整个标签集。我们可以做到这一点，而不会让函数的属性搞得一团糟：

class Selector(set):
    def tag(self, *tags):
        tags = set(tags)
        def decorator(func):
            if hasattr(func, "_tags"):
                func._tags.update(tags)
            else:
                func._tags = tags
            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                return func(*args, **kwargs) if self & func._tags else None
            wrapper._tags = func._tags
            return wrapper
        return decorator

这样做有一些好处——你可以查看这个函数所有的标签，还可以通过多个装饰器或者在一个装饰器中给出多个标签来进行标记。

@selector.tag('a')
@selector.tag('b')
def foo():
    ...


#Or, equivalently:
@selector.tag('a', 'b')
def foo():
    ...

使用 functools.wraps() 也意味着这个函数保持了它原来的“身份”（文档字符串、名称等等）。

编辑：如果你想要消除一些包装：

    def decorator(func):
        if hasattr(func, "_tagged_function"):
            func = func._tagged_function
        if hasattr(func, "_tags"):
            func._tags.update(tags)
        else:
            func._tags = tags
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs) if self & func._tags else None
        wrapper._tagged_function = func
        wrapper._tags = func._tags
        return wrapper