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)

Asked By: laridzhang

||

Source

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)

Answered By: Damian

Answer 2

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

Answered By: Nande

Can I pass parameters to the iter function in Python?

Question:

Answers:

Can I pass parameters to the __iter__ function in Python?

Question:

Answers:

Can I pass parameters to the iter function in Python?