Can I pass parameters to the __iter__ function in Python?

Question

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)

Answer 1

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)