CODEWITHSUNDEEP
python
laridzhang
laridzhang

Reputation: 1811

Can I pass parameters to the __iter__ function in Python?

I want to pass a parameter to control the size of yielded data. For example,

class Something:
    def __init__(self):
        self.all_data = list(range(23))

    def __iter__(self, size):
        random.shuffle(self.all_data)
        batch = list()
        for i in self.all_data:
            batch.append(i)
            if len(batch) >= size:
                yield batch
                batch = list()

a = Something()
for i in a:  # TypeError: __iter__() missing 1 required positional argument: 'size'
    print(i)
for i in a(5):  # TypeError: 'Something' object is not callable
    print(i)

Upvotes: 10

Views: 5261

Answers (2)

Nande
Nande

Reputation: 442

i think it would be clearer to create a function that returns an iterator. or pass it in init

def __init__(self, data):
    #your code

for i in Something(42): pass

Upvotes: 0

Damian
Damian

Reputation: 568

You can get your desirable effect by implementing __call__ method which will collect data for __iter__ to use, I set self.size in init constructor to ensure attribute existence.

It's not exactly what you want because it's not passing argument directly but via object's state, but it's working in case of for i in a(5)

please consider case when someone will use your code like this for i in a and by that, it can alter logic, you can check if self.size isn't None for example.

import random


class Something:
    def __init__(self):
        self.all_data = list(range(23))
        self.size = 0

    def __iter__(self):
        random.shuffle(self.all_data)
        batch = list()
        for i in self.all_data:
            batch.append(i)
            if len(batch) >= self.size:
                yield batch
                batch = list()

    def __call__(self, n):
        self.size = n
        return self


a = Something()

for i in a(0):
    print(i)

Upvotes: 13

Related Questions