__all__ = ['Iter']
from itertools import *
class Iter(object):
'''A wrapper class providing a rich-iterator API.
From the user point of view, this class supersedes the builtin iter()
function: like iter(), it is called as Iter(iterable) or (less frequently)
as Iter(callable,sentinel) and it returns an iterator. The returned
iterator, in addition to the basic iterator protocol, provides a rich API,
exposing as methods most functions of the itertools module. Notably, two
frequently used itertools functions, chain and islice, are conveniently
exposed as addition and slicing operators, respectively.
'''
def __init__(self, *args): self._it = iter(*args)
def __iter__(self): return self
def next(self): return self._it.next()
def __add__(self, other):
if not isinstance(other,Iter):
raise TypeError('can only add Iter (not "%s") to Iter' %
other.__class__.__name__)
return Iter(chain(self._it, other._it))
def __mul__(self, num): return Iter(chain(*tee(self._it,num)))
__rmul__ = __mul__
def __getitem__(self, index):
if isinstance(index, int):
try: return islice(self._it, index, index+1).next()
except StopIteration:
raise IndexError('Index %d out of range' % index)
else:
start,stop,step = index.start, index.stop, index.step
if start is None: start = 0
if step is None: step = 1
return Iter(islice(self._it, start, stop, step))
def enumerate(self): return Iter(enumerate(self._it))
def map(self, func): return Iter(imap(func,self))
def zip(self, *others): return Iter(izip(self._it, *others))
def filter(self, predicate): return Iter(ifilter(predicate,self._it))
def filterfalse(self, predicate): return Iter(ifilterfalse(predicate,self._it))
def cycle(self): return Iter(cycle(self._it))
def takewhile(self, predicate): return Iter(takewhile(predicate, self._it))
def dropwhile(self, predicate): return Iter(dropwhile(predicate, self._it))
def groupby(self, keyfunc=None): return Iter(groupby(self._it, keyfunc))
def copy(self):
self._it, new = tee(self._it)
return Iter(new)
def irange(*args):
'''Return an Iter-wrapped xrange object.'''
return Iter(xrange(*args))
if __name__ == '__main__':
# Example: A composite iterator over two files specified as follows:
# - each fetched line is right stripped.
# - the first 3 lines of the first file are fetched.
# - the first line of the second file is skipped and its next 4 lines are fetched.
# - empty lines (after the right stripping) are filtered out.
# - the remaining lines are enumerated.
import sys
f1,f2 = [Iter(open(f)).map(str.rstrip) for f in sys.argv[1:3]]
for i,line in (f1[:3] + f2[1:5]).filter(None).enumerate():
print i,repr(line)
