Search code examples
pythondecoratorpython-decorators

Decorator that will also define another function


I have a class that contains a list of "features" wherein each feature is the name of a function found within that class. It's a way of controlling which features are added to a timeseries dataframe dynamically. I have a function in another class which is meant to take an existing feature and retrieve its future state. So based on the way the current structure is set up, I need to dynamically add a function by appending "_after" to the end of its name.

I've got my decorator set up so far that the features list is updated with the new function name, but I don't know how to declare an additional function from within the decorator. Ultimately I don't even need to wrap the original function, I just need to create a new one using the naming convention of the old one.

In this contrived example, Dog should in the end have two functions: bark() and bark_after(). The Features list should contain ['bark', 'bark_again']. I'd also like to not have to pass Features in explicitly to the decorator.

class Dog:
    Features = ['bark']

    def __init__(self, name):
        self.name = name
        
    def gets_future_value(*args):
        def decorator(function):
            new_function = f'{function.__name__}_after'
            args[0].append(new_function)
            return function
        return decorator
    
    @gets_future_value(Features)
    def bark(self):
        print('Bark!')

d = Dog('Mr. Barkington')
print(d.Features)
d.bark()
d.bark_after()

Solution

  • I think what you need is a class decorator:

    import inspect
    
    
    def add_after_methods(cls):
        features = set(cls.Features)  # For fast membership testing.
        isfunction = inspect.isfunction
    
        def isfeature(member):
            """Return True if the member is a Python function and a class feature."""
            return isfunction(member) and member.__name__ in features
    
        # Create any needed _after functions.
        for name, member in inspect.getmembers(cls, isfeature):
            after_func_name = name + '_after'
            def after_func(*args, **kwargs):
                print(f'in {after_func_name}()')
            setattr(cls, after_func_name, after_func)
            cls.Features.append(after_func_name)
        return cls
    
    
    @add_after_methods
    class Dog:
        Features = ['bark']
    
        def __init__(self, name):
            self.name = name
    
        def bark(self):
            print('Bark!')
    
    
    d = Dog('Mr. Barkington')
    print(d.Features)  # -> ['bark', 'bark_after']
    d.bark()  # -> Bark!
    d.bark_after()  # -> in bark_after()