Welcome, guest | Sign In | My Account | Store | Cart
__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' % 
        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)
            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)