简体   繁体   中英

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)

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)

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

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM