Search code examples
pythondecoratorchaining

passing variables between two python decorators


Is there are a way to pass a variable between two python decorators applied to the same function? The goal is for one of the decorators to know that the other was also applied. I need something like decobar_present() from the example below:

def decobar(f):
    def wrap():
        return f() + "bar"
    return wrap

def decofu(f):
    def wrap():
        print decobar_present() # Tells me whether decobar was also applied
        return f() + "fu"
    return wrap

@decofu
@decobar
def important_task():
    return "abc"

More generally I would like to be able to modify the behavior of decofu depending on whether decobar was also applied.


Solution

  • You can add the function to a "registry" when decobar is applied to it, then later check the registry to determine whether decobar was applied to the function or not. This approach requires preserving original function's __module__ and __name__ properties intact (use functools.wraps over the wrapper function for that).

    import functools
    
    class decobar(object):
        registry = set()
    
        @classmethod
        def _func_key(cls, f):
            return '.'.join((f.__module__, f.func_name))
    
        @classmethod
        def present(cls, f):
            return cls._func_key(f) in cls.registry
    
        def __call__(self, f):
            self.registry.add(self._func_key(f))
    
            @functools.wraps(f)
            def wrap():
                return f() + "bar"
            return wrap
    
    # Make the decorator singleton
    decobar = decobar()
    
    def decofu(f):
        @functools.wraps(f)
        def wrap():
            print decobar.present(f) # Tells me whether decobar was also applied
            return f() + "fu"
        return wrap
    
    @decofu
    @decobar
    def important_task():
        return "abc"
    

    Used a class to implement decobar, as it keeps registry and present() in a single namespace (which feels slighly cleaner, IMO)