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

时间:2013-01-28 18:16:10

标签: python decorator

我正在创建一个功能标记系统,以启用或禁用基于标记的功能:

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

编辑:添加了装饰者返回。

2 个答案:

答案 0 :(得分:4)

问题在于,如果你装饰它两次,一个返回函数,一个不返回任何内容。

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

答案 1 :(得分: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"