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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
|
#!/usr/bin/env python3
# © Copyright 2021-2022, Scott Gasch
"""A collection if :class:`Iterator` subclasses that can be composed
with another iterator and provide extra functionality. e.g.
+ :class:`PeekingIterator`
+ :class:`PushbackIterator`
+ :class:`SamplingIterator`
"""
import random
from collections.abc import Iterator
from typing import Any, List, Optional
class PeekingIterator(Iterator):
"""An iterator that lets you :method:`peek` at the next item on deck.
Returns None when there is no next item (i.e. when
:method:`__next__` will produce a StopIteration exception).
>>> p = PeekingIterator(iter(range(3)))
>>> p.__next__()
0
>>> p.peek()
1
>>> p.peek()
1
>>> p.__next__()
1
>>> p.__next__()
2
>>> p.peek() == None
True
>>> p.__next__()
Traceback (most recent call last):
...
StopIteration
"""
def __init__(self, source_iter: Iterator):
self.source_iter = source_iter
self.on_deck: List[Any] = []
def __iter__(self) -> Iterator:
return self
def __next__(self) -> Any:
if len(self.on_deck) > 0:
return self.on_deck.pop()
else:
item = self.source_iter.__next__()
return item
def peek(self) -> Optional[Any]:
if len(self.on_deck) > 0:
return self.on_deck[0]
try:
item = self.source_iter.__next__()
self.on_deck.append(item)
return self.peek()
except StopIteration:
return None
class PushbackIterator(Iterator):
"""An iterator that allows you to push items back
onto the front of the sequence. e.g.
>>> i = PushbackIterator(iter(range(3)))
>>> i.__next__()
0
>>> i.push_back(99)
>>> i.push_back(98)
>>> i.__next__()
98
>>> i.__next__()
99
>>> i.__next__()
1
>>> i.__next__()
2
>>> i.push_back(100)
>>> i.__next__()
100
>>> i.__next__()
Traceback (most recent call last):
...
StopIteration
"""
def __init__(self, source_iter: Iterator):
self.source_iter = source_iter
self.pushed_back: List[Any] = []
def __iter__(self) -> Iterator:
return self
def __next__(self) -> Any:
if len(self.pushed_back) > 0:
return self.pushed_back.pop()
return self.source_iter.__next__()
def push_back(self, item: Any):
self.pushed_back.append(item)
class SamplingIterator(Iterator):
"""An iterator that simply echoes what source_iter produces but also
collects a random sample (of size sample_size) of the stream that can
be queried at any time.
Note that until sample_size elements have been seen the sample will
be less than sample_size elements in length.
Note that if sample_size is > len(source_iter) then it will produce
a copy of source_iter.
>>> import collections
>>> import random
>>> random.seed(22)
>>> s = SamplingIterator(iter(range(100)), 10)
>>> s.__next__()
0
>>> s.__next__()
1
>>> s()
[0, 1]
>>> collections.deque(s)
deque([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])
>>> s()
[78, 18, 47, 83, 93, 26, 25, 73, 94, 38]
"""
def __init__(self, source_iter: Iterator, sample_size: int):
self.source_iter = source_iter
self.sample_size = sample_size
self.resovoir: List[Any] = []
self.stream_length_so_far = 0
def __iter__(self) -> Iterator:
return self
def __next__(self) -> Any:
item = self.source_iter.__next__()
self.stream_length_so_far += 1
# Filling the resovoir
pop = len(self.resovoir)
if pop < self.sample_size:
self.resovoir.append(item)
if self.sample_size == (pop + 1): # just finished filling...
random.shuffle(self.resovoir)
# Swap this item for one in the resovoir with probabilty
# sample_size / stream_length_so_far. See:
#
# https://en.wikipedia.org/wiki/Reservoir_sampling
else:
r = random.randint(0, self.stream_length_so_far)
if r < self.sample_size:
self.resovoir[r] = item
return item
def __call__(self) -> List[Any]:
return self.resovoir
if __name__ == '__main__':
import doctest
doctest.testmod()
|
