class IterN (object):
def __init__ (self, iterable, n, end_marker=None):
if n == 1:
raise TypeError ("n = 1? woot?")
self.iterator = iter (iterable)
self.n = n
self.end_marker = end_marker
self.finished = False
def __iter__ (self):
return self
def next (self):
if self.finished:
raise StopIteration ()
res = []
i = 0
try:
# pull n items from the iterator
for i in range (self.n):
res.append (self.iterator.next ())
# set the first-invalid-index indicator to one past the length
i += 1
except StopIteration, x:
if i == 0:
# we couldn't get anything from the iterator
raise StopIteration ()
# we got less than expected, fix the remaining with end markers and
# make sure to stop the next time
res.extend ([self.end_marker] * (self.n - i))
self.finished = True
return res, i
def main ():
for seq, n in IterN ([1, 2, 3, 4, 5], 2):
if n < 2:
break
a, b = seq
# do something useful with a and b
print a, b
if __name__ == '__main__':
main ()