The decorator functions as tags as callables
I am creating a feature tag system to enable or disable tag-based features:
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 is not called. I was wondering if there is a way to make it get called if any of the tags are present in the selector. Ideally, the solution makes it so that tags are checked only once, rather than every time the function is called.
Note: I do this to select which tests to run based on environment variables in unit tests unittest
. My actual implementation uses unittest.skip
.
EDIT: Added decorator return.
source to share
The problem is, if you decorate it twice, a function is returned that doesn't return anything.
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 that in any order, you always get nothing. What you need to do is save a set of tags for each object and check that the whole set is 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 is possible to check all the tags that a function has, and you can mark them with multiple decorators or give many tags in one decorator.
@selector.tag('a')
@selector.tag('b')
def foo():
...
#Or, equivalently:
@selector.tag('a', 'b')
def foo():
...
Usage functools.wraps()
also means that the function retains its original identity (docstrings, name, etc.).
Edit: If you want to make some shell exceptions:
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
source to share