Class Iter

recipes/Python/498272_Rich_iterator_wrapper/recipe-498272.py:5–52

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 additi

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):