zezollo
zezollo

Reputation: 5017

Creating an iterator and choose among several sources of iteration

I want to create an Iterator that would have several "sources" of objects to iterate. I would like to be able to give its __next__() method an optional keyword argument that would offer the possibility to choose the source (no keyword argument would mean, just choose a source randomly).

Using __next__() causes problems (see below), so as a workaround I've written this unsatisfying code:

#!/usr/bin/env python3

import random

class infinite_iterator(object):

    def __init__(self, sources):
        self.collector = []
        for s in sources:
            random.shuffle(s)
            self.collector.append([])
        self.sources = sources

    def __iter__(self):
        return self

    # Workaround: calling it next instead of __next__
    # (not the python3 way...)
    def next(self, **kwargs):
        sce_nb = random.choice([ n for n in range(len(self.sources)) ])
        if 'choice' in kwargs:
            sce_nb = kwargs['choice']
        self.collector[sce_nb].append(self.sources[sce_nb][0])
        output = self.sources[sce_nb].pop(0)
        # Repopulate any empty 'source'
        if not self.sources[sce_nb]:
            (self.sources[sce_nb], self.collector[sce_nb]) = \
                         (self.collector[sce_nb], self.sources[sce_nb])
            random.shuffle(self.sources[sce_nb])
        return output


S = infinite_iterator([["Adam", "Paul", "Oliver", "Aaron", "Joshua", "Jack"],
                       ["Mary", "Sophia", "Emily", "Isobel", "Grace", "Alice", "Lucy"]])

print("Any name: " + S.next())
print("Any girl's name: " + S.next(choice=1))
print("Any boy's name: " + S.next(choice=0))

Problem is, if I want to write def __next__(self, **kwargs): to make infinite_iterator a real Iterator, then of course I want to write:

print("Any name: " + next(S))
print("Any girl's name: " + next(S, choice=1))
print("Any boy's name: " + next(S, choice=0))

but get an error (2d line):

TypeError: next() takes no keyword arguments

I thought this call next(S, choice=1) would use the __next__() function defined in the object. Because of this error, I think that on one hand it actually does not. This could be expectable because it's not exactly a redefinition, as infinite_iterator does not inherit from an "Iterator object" (as far as I understand, there's no such object). But on another hand, if I call only next(S) it works, and in this case, my __next__() method is really called (it chooses randomly a list to iterate over).

Finally, is there a solution to realize such an Iterator?

Upvotes: 2

Views: 111

Answers (1)

tzaman
tzaman

Reputation: 47770

Write a generator (technically a coroutine) and use send to pass in your choices:

import random
def selector(sources):
    option = None
    while True:
        if option is None:
            option = yield random.choice(random.choice(sources))
        else:
            option = yield random.choice(sources[option])

s = selector([["Adam", "Paul", "Oliver", "Aaron", "Joshua", "Jack"],
              ["Mary", "Sophia", "Emily", "Isobel", "Grace", "Alice", "Lucy"]])

print("Any name: " + next(s))
print("Any girl's name: " + s.send(1))
print("Any boy's name: " + s.send(0))

The way this works is every time you call send, that value gets passed in to the option variable which is used to make the selection for the next iteration of the while loop.

You'll have to start it up by calling next at least once so it hits the first yield statement, but after that it'll work both ways (use next or send(None) for a full random choice, or send(option) for a specific source.

Upvotes: 0

Related Questions