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

Updated: Use itertools.groupby instead. See my comment below.

This is a variation of itertools.groupby. The itertools.groupby iterator assumes that the input is not sorted but will fit in memory. This iterator has the same API, but assumes the opposite.

Python, 117 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
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
__docformat__ = "restructuredtext"


class peekable:

    """Make an iterator peekable.

    This is implemented with an eye toward simplicity.  On the downside,
    you can't do things like peek more than one item ahead in the
    iterator.  On the bright side, it doesn't require anything from
    itertools, etc., so it's less likely to encounter strange bugs,
    which occassionally do happen.

    Example usage::

        >>> numbers = peekable(range(6))
        >>> numbers.next()
        0
        >>> numbers.next()
        1
        >>> numbers.peek()
        2
        >>> numbers.next()
        2
        >>> numbers.next()
        3
        >>> for i in numbers:
        ...     print i
        ... 
        4
        5

    """

    _None = ()  # Perhaps None is a valid value.

    def __init__(self, iterable):
        self._iterable = iter(iterable)
        self._buf = self._None

    def __iter__(self):
        return self

    def _is_empty(self):
        return self._buf is self._None

    def peek(self):
        """Peek at the next element.

        This may raise StopIteration.

        """
        if self._is_empty():
            self._buf = self._iterable.next()
        return self._buf

    def next(self):
        if self._is_empty():
            return self._iterable.next()
        ret = self._buf
        self._buf = self._None
        return ret


def groupbysorted(iterable, keyfunc=None):

    """This is a variation of itertools.groupby.

    The itertools.groupby iterator assumes that the input is not sorted
    but will fit in memory.  This iterator has the same API, but assumes
    the opposite.

    Example usage::

        >>> for (key, subiter) in groupbysorted(
        ...     ((1, 1), (1, 2), (2, 1), (2, 3), (2, 9)), 
        ...     keyfunc=lambda row: row[0]):
        ...     print "New key:", key
        ...     for x in subiter:
        ...         print "Row:", x
        ... 
        New key: 1
        Row: (1, 1)
        Row: (1, 2)
        New key: 2
        Row: (2, 1)
        Row: (2, 3)
        Row: (2, 9)

    This requires the peekable class.  See my comment here_.

    Note, you must completely iterate over each subiter or groupbysorted will
    get confused.

    .. _here:
       http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/304373

    """

    iterable = peekable(iterable)

    if not keyfunc:
        def keyfunc(x):
            return x

    def peekkey():
        return keyfunc(iterable.peek())

    def subiter():
        while True:
            if peekkey() != currkey:
                break
            yield iterable.next()

    while True:
        currkey = peekkey()
        yield (currkey, subiter())

I wrote this recipe out of my own personal need, of course. I was using an iterator to wrap a cursor that was returning a huge data set from a database. I first tried to use itertools.groupby. That failed. Hence, I wrote this. I'm relying on the database to sort the records.

Here's a big shout out to Danny Yoo (a fellow member of BayPiggies) and the guy (unknown) who wrote the peek recipe. Thanks!

1 comment

Shannon -jj Behrens (author) 13 years, 3 months ago  # | flag

It turns out I was wrong about itertools.groupby. It works the same as this code does, so you should use that instead.

Created by Shannon -jj Behrens on Tue, 16 Jan 2007 (PSF)
Python recipes (4591)
Shannon -jj Behrens's recipes (19)

Required Modules

  • (none specified)

Other Information and Tasks