Search code examples
pythondecorator

Iterable decorator


I have a class in python with the following structure:

class MyClass:
    def __init__(self, iterable):
        self.iterable = iterable

    def example_0(self):
        for i in self.iterable:
            print(i)

    def example_1(self):
        for i in self.iterable:
            print(i + 1)

    def example_2(self):
        for i in self.iterable:
            print(i + 2)

That is, I have several methods that run different operations on an iterable that is an attribute of the class. I need to run for i in self.iterable for each method in the class, and I would otherwise like to use a decorator to all these methods, something like:

class MyClass:
    def __init__(self, iterable):
        self.iterable = iterable

    @iterate
    def example_0(self, i):
        print(i)

    @iterate
    def example_1(self, i):
        print(i + 1)

    @iterate
    def example_2(self, i):
        print(i + 2)

Can you help me write this decorator such that the behavior of my class is the same as the new class?


I tried this:

class MyClass:
    def __init__(self, iterable):
        self.iterable = iterable

    def iterate(self, f):
        def func(*args, **kwargs):
            for i in self.iterable:
                f(self, i, *args, **kwargs)
        return func

    @iterate
    def example_0(self, i):
        print(i)

    @iterate
    def example_1(self, i):
        print(i + 1)

    @iterate
    def example_2(self, i):
        print(i + 2)

and it returns: TypeError: iterate() missing 1 required positional argument: 'f'.

My main issue is I'm not sure how to put the decorator within my class, as it's iterating over an attribute of the class.


Solution

  • I don't think this has much advantage, but here's an option. This also allows for additional parameters.

    from functools import wraps
    
    def iterated(f):
        @wraps(f)
        def _f(self, *args, **kwargs):
            for i in self.iterable:
                f(self, i, *args, **kwargs)
        return _f
    

    Example:

    In [11]: class MyClass2:
        ...:     def __init__(self, iterable):
        ...:         self.iterable = iterable
        ...: 
        ...:     @iterated
        ...:     def example_0(self, i, k):
        ...:         print(i + k)
    
    In [12]: MyClass2([1,2,3]).example_0(2)
    3
    4
    5
    

    This is to be defined outside the class. In your failing implementation, you forgot the self parameter in the returned closure.