File: grouping.py

package info (click to toggle)
python-picklable-itertools 0.1.1-2
  • links: PTS, VCS
  • area: main
  • in suites: buster
  • size: 172 kB
  • sloc: python: 1,222; makefile: 3
file content (65 lines) | stat: -rw-r--r-- 2,161 bytes parent folder | download | duplicates (3)
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
from .base import BaseItertool
from .iter_dispatch import iter_


class _grouper(BaseItertool):
    def __init__(self, value, iterator, groupby_obj):
        self._value = value
        self._groupby = groupby_obj
        self._key = self._groupby.key(self._value)
        self._initialized = False
        self._iterator = iterator
        self.stream_ended = False
        self._done = False

    def __next__(self):
        if not self._initialized:
            self._initialized = True
            return self._value
        else:
            if self._done:
                raise StopIteration
            try:
                value = next(self._iterator)
            except StopIteration:
                self.stream_ended = True
                self._done = True
                raise
            if self._groupby.key(value) != self._key:
                self.terminal_value = value
                self._done = True
                raise StopIteration
            return value


class groupby(BaseItertool):
    """groupby(iterable[, keyfunc]) -> create an iterator which returns
    (key, sub-iterator) grouped by each value of key(value).
    """
    def __init__(self, iterable, key=None):
        self._keyfunc = key
        self._iterator = iter_(iterable)
        self._current_key = self._initial_key = object()

    def key(self, value):
        if self._keyfunc is None:
            return value
        else:
            return self._keyfunc(value)

    def __next__(self):
        if not hasattr(self, '_current_grouper'):
            value = next(self._iterator)
            self._current_grouper = _grouper(value, self._iterator, self)
            return self.key(value), self._current_grouper
        else:
            while True:
                try:
                    next(self._current_grouper)
                except StopIteration:
                    break
            if self._current_grouper.stream_ended:
                raise StopIteration
            value = self._current_grouper.terminal_value
            self._current_grouper = _grouper(value, self._iterator, self)
            return self.key(value), self._current_grouper