Welcome, guest | Sign In | My Account | Store | Cart

This recipe may be of interest to those who make heavy use of the itertools module. It provides a wrapper class that exposes most itertools functions as methods, plus a few more. Moreover, two frequently used itertools functions, chain and islice, are conveniently exposed as addition and slicing operators, respectively.

Python, 72 lines
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
__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)

Iterables, iterators, generators, generator comprehensions are some of the core python concepts, used extensively both for their elegance and efficiency. The itertools module provides several basic iterator building blocks by composing and transforming one or more existing iterables.

This recipe can be seen as an object-oriented API to itertools utilities. The Iter class wraps an arbitrary iterable into a "rich iterator" object, an iterator that provides equivalent methods to most itertools functions. Two functions in particular, chain and islice, are exposed as addition and slicing operators, which results in terse and yet more readable (IMHO) code.

1 comment

Christopher Dunn 17 years, 2 months ago  # | flag

Interesting. This makes Python look more like Ruby.