Question

I’m creating a function tagging system, to enable or disable functions based on tags:

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'

My question is about the last function, foo3. I understand why it isn’t being called. I was wondering if there’s a way to make it so that it is called if any of the tags are present in the selector. Ideally, the solution makes it so the tags are only checked once, not every time the function is called.

A side note: I’m doing this to select tests to run based on environment variables in unittest unit tests. My actual implementation uses unittest.skip.

EDIT: Added the decorator return.

Answer 1

The issue is that if you decorate it twice, one returns the function, one returns nothing.

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

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

This means, in whatever order, you will always get nothing. What you need to do is keep a set of tags on each object, and check that whole set at once. We can do this nicely without polluting namespaces with function attributes:

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

This gives some bonuses – it’s possible to inspect all the tags the function has, and it’s possible to tag with multiple decorators or by giving many tags in a single decorator.

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


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

The use of functools.wraps() also means the function keeps it’s original ‘identity’ (docstrings, name, etc…).

Edit: If you wanted to do some wrapper elimination:

    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