I use iter
with two arguments and was wondering if there is an equivalent which would accept a more complex sentinel?
As an example, in the code below
# returns digits 1 to 10 upon subsequently calling .create(),
# then returns 'END' each time afterwards
class MyClass:
def __init__(self):
self.counter = 0
def create(self):
self.counter += 1
if self.counter > 10:
return 'END'
else:
return self.counter
c = MyClass()
for i in iter(c.create, 'END'):
print(i)
the iteration ends upon getting 'END'
. I would like to have it end after getting, say, a total of two 'END'
(not necessarily one after the other - the code above will generate only 'END'
after the first ten calls but one can imagine a case where they are interlaced with other values).
In essence, I am looking for a way to use a more complex sentinel. Is there such a concept?
I know I can resort to the obvious-but-ugly code mentioned in a previous question, hovewer @HappyLeapSecond answer is so elegant I would like to keep the spirit (I went though itertools
but none of the available methods seem to do the job)
You could use itertools.takewhile
with a stateful predicate
. For example:
>>> from itertools import takewhile
>>> def create_predicate(func, max_count):
"""Return False only when func evaluates True for the max_count time."""
def predicate(elem):
if func(elem):
predicate.count += 1
if predicate.count == max_count:
return False
return True
predicate.count = 0
return predicate
>>> list(takewhile(create_predicate(lambda elem: elem % 3 == 0, 3), range(1, 20)))
[1, 2, 3, 4, 5, 6, 7, 8]
In the example above, create_predicate(lambda elem: elem % 3 == 0, 3)
creates a complex predicate function that will stop the iteration on the third multiple of three. In your case,
I would like to have it end after getting, say, a total of two
'END'
you could use create_predicate(lambda elem: elem == 'END', 2)
.