File: iter_dispatch.py

package info (click to toggle)
python-picklable-itertools 0.1.1-2
  • links: PTS, VCS
  • area: main
  • in suites:
  • size: 172 kB
  • sloc: python: 1,222; makefile: 3
file content (93 lines) | stat: -rw-r--r-- 2,613 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
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
import io

import six
try:
    import numpy
    NUMPY_AVAILABLE = True
except ImportError:
    numpy = None
    NUMPY_AVAILABLE = False


from .base import BaseItertool


def iter_(obj):
    """A custom replacement for iter(), dispatching a few custom picklable
    iterators for known types.
    """
    if six.PY2:
        file_types = file,  # noqa
    if six.PY3:
        file_types = io.IOBase,
        dict_items = {}.items().__class__
        dict_values = {}.values().__class__
        dict_keys = {}.keys().__class__
        dict_view = (dict_items, dict_values, dict_keys)

    if isinstance(obj, dict):
        return ordered_sequence_iterator(list(obj.keys()))
    if isinstance(obj, file_types):
        return file_iterator(obj)
    if six.PY2:
        if isinstance(obj, (list, tuple)):
            return ordered_sequence_iterator(obj)
        if isinstance(obj, xrange):  # noqa
            return range_iterator(obj)
        if NUMPY_AVAILABLE and isinstance(obj, numpy.ndarray):
            return ordered_sequence_iterator(obj)
    if six.PY3 and isinstance(obj, dict_view):
        return ordered_sequence_iterator(list(obj))
    return iter(obj)


class range_iterator(BaseItertool):
    """A picklable range iterator for Python 2."""
    def __init__(self, xrange_):
        self._start, self._stop, self._step = xrange_.__reduce__()[1]
        self._n = self._start

    def __next__(self):
        if (self._step > 0 and self._n < self._stop or
                self._step < 0 and self._n > self._stop):
            value = self._n
            self._n += self._step
            return value
        else:
            raise StopIteration


class file_iterator(BaseItertool):
    """A picklable file iterator."""
    def __init__(self, f):
        self._f = f

    def __next__(self):
        line = self._f.readline()
        if not line:
            raise StopIteration
        return line

    def __getstate__(self):
        name, pos, mode = self._f.name, self._f.tell(), self._f.mode
        return name, pos, mode

    def __setstate__(self, state):
        name, pos, mode = state
        self._f = open(name, mode=mode)
        self._f.seek(pos)


class ordered_sequence_iterator(BaseItertool):
    """A picklable replacement for list and tuple iterators."""
    def __init__(self, sequence):
        self._sequence = sequence
        self._position = 0

    def __next__(self):
        if self._position < len(self._sequence):
            value = self._sequence[self._position]
            self._position += 1
            return value
        else:
            raise StopIteration