CODEWITHSUNDEEP
pythondecoratorpython-decorators
SnakeWasTheNameTheyGaveMe
SnakeWasTheNameTheyGaveMe

Reputation: 2562

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()

Upvotes: 1

Views: 657

Answers (1)

martineau
martineau

Reputation: 123453

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()

Upvotes: 1

Related Questions