Coverage for pycommons/ds/sequences.py: 97%
66 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-15 00:53 +0000
1"""Tools for working with sequences."""
3from heapq import merge
4from itertools import chain
5from typing import Callable, Final, Generator, Iterable, Iterator, TypeVar
7from pycommons.types import type_error
9#: the type of the element of the sequences to process
10T = TypeVar("T")
13class __Reiterator(Iterable[T]):
14 """The internal class for re-iteration."""
16 def __init__(self, source: Iterator[T], can_unpack: bool) -> None:
17 """
18 Create the re-iterator.
20 :param source: the source
21 """
22 #: the original source
23 self.__source: Final[Iterator[T]] = source
24 #: we store all elements from source for re-iteration
25 self.__more: Final[list[T]] = []
26 #: the unpacking state
27 self._unpack_state: int = 1 if can_unpack else 0
29 def __iter__(self) -> Generator[T, None, None]:
30 """
31 Generate the sequence of elements.
33 :returns: the generator
34 """
35 if self._unpack_state > 1:
36 raise ValueError("Reiterator has been unpacked!")
37 get_length: Callable[[], int] = self.__more.__len__
38 get: Callable[[int], T] = self.__more.__getitem__
39 append: Final[Callable[[T], None]] = self.__more.append
40 nexter: Final[Callable[[], T]] = self.__source.__next__
41 pos: int = 0 # The next position in __more.
42 skip: int = -1 # The last element in __more we took from __source.
43 try: # noqa: PLW0717 # We always get to StopIteration of __source.
44 while True: # Until we reached the end of list and end of iter.
45 ll: int = get_length() # Get length (may have changed).
46 while pos < ll: # First, return all elements from __more.
47 if pos != skip: # Skip = the element we returned last.
48 yield get(pos) # Yield the __more list element.
49 pos += 1 # Increment position.
50 skip = get_length() # Get the length: May have changed!
51 if skip > ll: # Maybe something changed since yield.
52 skip = -1 # Invalidate skip, we instead return from list.
53 continue
54 other = nexter() # Get next actual iteration element.
55 append(other) # Store the element in the internal list.
56 yield other # Yield the element
57 except StopIteration:
58 pass
60 def _unpack(self) -> Iterator[T]:
61 """
62 Unpack the re-iterable.
64 :return: the unpacked re-iterable.
65 """
66 if self._unpack_state <= 0:
67 raise ValueError("This re-iterable cannot be unpacked.")
68 if self._unpack_state > 1:
69 raise ValueError("Re-iterable has already been unpacked!")
70 self._unpack_state = 2
71 return chain(self.__more, self.__source) if self.__more \
72 else self.__source
75def reiterable(source: Iterable[T] | Iterator[T],
76 unpackable: bool = False) -> Iterable[T]:
77 """
78 Ensure that an :class:`Iterable` can be iterated over multiple times.
80 This function will solidify an :class:`Iterator` into an
81 :class:`Iterable`. In Python, :class:`Iterator` is a sub-class of
82 :class:`Iterable`. This means that if your function accepts instances of
83 :class:`Iterable` as input, it may expect to be able to iterate over them
84 multiple times. However, if an :class:`Iterator` is passed in, which also
85 is an instance of :class:`Iterable` and thus fulfills the function's type
86 requirement, this is not the case. A typical example of this would be if
87 a :class:`Generator` is passed in. A :class:`Generator` is an instance of
88 :class:`Iterator`, which, in turn, is an instance of :class:`Iterable`.
89 However, you can iterate over a :class:`Generator` only once.
91 For such single-use objects, a new :class:`Iterable` wrapper is created.
92 This wrapper will iterate over the original sequence, but cache all
93 elements in an internal :class:`list`. When you iterate over the sequence
94 again, the elements in the :class:`list` will be used. This means that
95 all elements of the original sequence will be stored in memory. However,
96 they are only stored if/when they are actually accessed via the iteration
97 sequence. If you do not iterate over them completely, they are not all
98 stored.
100 This form of re-iterabling is useful if you maybe generate items from a
101 slower sequence or do not plan to use all of them. If you want to use
102 all elements several times anyway, it may be more efficient to just wrap
103 the original source into a :class:`tuple`. But if, for example, your
104 sequence is the result of iterating over a directory tree on the file
105 system, or maybe if it comes from loading a file, then using
106 :func:`reiterable` could be useful.
108 This is also true if you actually process the generated sequence in some
109 way that may fail or terminate early. Then, first loading all data into
110 a :class:`tuple` may be annoying if your first processed element after
111 that causes a failure or early termination. The bit of overhead of
112 :func:`reiterable` may then well be worth your while.
114 Of course, this can only work if the :class:`Iterator` is not otherwise
115 used after calling this function. If you extract elements from the
116 :class:`Iterator` by yourself otherwise, maybe via :func:`next`, then
117 :func:`reiterable` cannot work. However, if you only apply :func:`next`
118 or other looping paradigms to the :class:`Iterable` returned by
119 :func:`reiterable`, then you can iterate as often as you want over a
120 :class:`Generator`, for example.
122 The function :func:`unpack_reiterable` can be used to "unpack" the
123 reiterable, returning a single-shot iterator of all the data.
124 The use case is as follows: Our CSV output facility in
125 :mod:`pycommons.io.csv` allows you to turn a sequence of arbitrary
126 data elements to a CSV text stream.
127 This may require two passes over the data:
128 In the first pass, one *may* determine the columns that are needed.
129 In some cases, we already know the columns and do not need to pass
130 over the data at all.
131 In some other cases, we may need one potentially incomplete pass over
132 the data, just to check if certain columns are needed.
133 For example, when writing out statistical data, we may first want to
134 check if there is any record where the geometric mean is defined.
135 Once we see the first such record, we know that we need to the
136 corresponding column and do not need to check any subsequent elements.
137 Then we can write the CSV header with the column names and start the
138 actual iteration.
139 We thus have one incomplete pass over the data (to find the columns)
140 and one complete pass (to write the data).
142 During re-iteration, we need to cache all already visited elements.
143 Thus, without unpacking, we would necessarily have to cache all the
144 elements of the sequence in memory, even if we only need to see them
145 once.
146 However, what we actually do is that we initialize the columns during
147 re-iteration, which may require no pass over the data (no caching), an
148 incomplete pass (thus only caching the actually visited elements), or
149 a complete pass (caching all elements).
150 Then, we unpack the re-iterable for the pass that actually writes the
151 data.
152 The resulting :class:`Iterable` will first iterate over the cached
153 elements and then over the remainder in the original :class:`Iterator`.
154 It thus requires the minimum amount of caching.
156 For enabling unpacking, you must set `unpackable` to `True`.
158 Notice that wrapping pure instances of :class:`Iterable` into re-iterables
159 has no effect and neither does unpacking.
160 The reason is that if the generation of arbitrarily many iterators is
161 already supported anyway, then we do not need to add the ability to do
162 this via the caching mechanism offered here.
164 :param source: the data source
165 :param unpackable: should the iterator be potentially unpackable? (
166 see :func:`unpack_reiterable`)
167 :returns: the resulting re-iterable iterator
168 :raises TypeError: if `source` is neither an :class:`Iterable` nor an
169 :class:`Iterator` or if `unpackable` is not a `bool`.
170 :raises ValueError: if `source` is a re-iterable that has been unpacked
171 via :func:`unpack_reiterable`.
173 >>> ri = reiterable(iter(range(4)))
174 >>> for i in ri:
175 ... print(i)
176 0
177 1
178 2
179 3
180 >>> for i in ri:
181 ... print(i)
182 0
183 1
184 2
185 3
187 >>> g = (i ** 2 for i in range(5))
188 >>> r = reiterable(g)
189 >>> tuple(r)
190 (0, 1, 4, 9, 16)
191 >>> tuple(r)
192 (0, 1, 4, 9, 16)
193 >>> tuple(r)
194 (0, 1, 4, 9, 16)
195 >>> tuple(r)
196 (0, 1, 4, 9, 16)
197 >>> tuple(g)
198 ()
200 >>> g = (i ** 2 for i in range(5))
201 >>> r = reiterable(g)
202 >>> i1 = iter(r)
203 >>> i2 = iter(r)
204 >>> next(i1)
205 0
206 >>> next(i2)
207 0
208 >>> next(i2)
209 1
210 >>> next(i1)
211 1
212 >>> next(i1)
213 4
214 >>> next(i1)
215 9
216 >>> next(i2)
217 4
218 >>> next(i2)
219 9
220 >>> i3 = iter(r)
221 >>> next(i3)
222 0
223 >>> next(i3)
224 1
225 >>> next(i3)
226 4
227 >>> next(i3)
228 9
229 >>> next(i3)
230 16
231 >>> next(i2)
232 16
233 >>> try:
234 ... next(i2)
235 ... except StopIteration as si:
236 ... print(type(si))
237 <class 'StopIteration'>
238 >>> try:
239 ... next(i3)
240 ... except StopIteration as si:
241 ... print(type(si))
242 <class 'StopIteration'>
243 >>> next(i1)
244 16
245 >>> try:
246 ... next(i1)
247 ... except StopIteration as si:
248 ... print(type(si))
249 <class 'StopIteration'>
251 >>> a = [1, 2, 3]
252 >>> reiterable(a) is a
253 True
255 >>> a = (1, 2, 3)
256 >>> reiterable(a) is a
257 True
259 >>> a = {1, 2, 3}
260 >>> reiterable(a) is a
261 True
263 >>> a = {1: 1, 2: 2, 3: 3}
264 >>> reiterable(a) is a
265 True
267 >>> k = a.keys()
268 >>> reiterable(k) is k
269 True
271 >>> k = a.values()
272 >>> reiterable(k) is k
273 True
275 >>> tuple(reiterable((x for x in range(5))))
276 (0, 1, 2, 3, 4)
278 >>> try:
279 ... reiterable(None)
280 ... except TypeError as te:
281 ... print(str(te)[:60])
282 source should be an instance of any in {typing.Iterable, typ
284 >>> try:
285 ... reiterable(1)
286 ... except TypeError as te:
287 ... print(str(te)[:60])
288 source should be an instance of any in {typing.Iterable, typ
290 >>> try:
291 ... reiterable((2, 3), "x")
292 ... except TypeError as te:
293 ... print(str(te)[:59])
294 unpackable should be an instance of bool but is str, namely
296 >>> type(merge_sorted_and_return_unique([1, 2, 3,], [2, 2]))
297 <class 'generator'>
299 >>> type(reiterable(merge_sorted_and_return_unique([1, 2, 3,], [2, 2])))
300 <class 'pycommons.ds.sequences.__Reiterator'>
302 >>> def myfunc():
303 ... if hasattr(myfunc, "a"):
304 ... s = getattr(myfunc, "a")
305 ... else:
306 ... s = 0
307 ... for i in range(s, s + 3):
308 ... yield i
309 ... setattr(myfunc, "a", s + 3)
310 >>> ri = reiterable(myfunc())
311 >>> tuple(ri)
312 (0, 1, 2)
313 >>> tuple(ri)
314 (0, 1, 2)
315 >>> ri = reiterable(myfunc())
316 >>> tuple(ri)
317 (3, 4, 5)
318 >>> tuple(ri)
319 (3, 4, 5)
320 """
321 if not isinstance(unpackable, bool):
322 raise type_error(unpackable, "unpackable", bool)
323 is_reiterator: Final[bool] = isinstance(source, __Reiterator)
324 if is_reiterator:
325 if source._unpack_state == 0: # type: ignore
326 return source
327 if source._unpack_state == 2: # type: ignore
328 raise ValueError("Unsafe operation: The source is a re-"
329 "iterable that has been unpacked.")
330 if is_reiterator or isinstance(source, Iterator):
331 return __Reiterator(source, unpackable) # type: ignore
332 if not isinstance(source, Iterable):
333 raise type_error(source, "source", (Iterable, Iterator))
334 return source # other iterables can be returned as-is
337def unpack_reiterable(source: Iterable[T]) -> Iterable[T]:
338 """
339 Unpack a re-iterable sequence.
341 See the documentation of :func:`reiterable` for more details.
343 :param source: the source
344 :return: the unpacked sequence
345 :raises ValueError: if the source is a re-iterable that was created
346 using :func:`reiterable` with `unpackable` set to `False` or if
347 the re-iterable has already been unpacked.
349 >>> x = reiterable(iter(range(5)), unpackable=True)
350 >>> y = iter(x)
351 >>> next(y)
352 0
354 >>> next(y)
355 1
357 >>> z = unpack_reiterable(x)
358 >>> for i in z:
359 ... print(i)
360 0
361 1
362 2
363 3
364 4
366 >>> try:
367 ... unpack_reiterable(x)
368 ... except ValueError as ve:
369 ... print(ve)
370 Re-iterable has already been unpacked!
372 >>> try:
373 ... reiterable(x, False)
374 ... except ValueError as ve:
375 ... print(ve)
376 Unsafe operation: The source is a re-iterable that has been unpacked.
378 >>> try:
379 ... unpack_reiterable(reiterable(iter(range(5)), False))
380 ... except ValueError as ve:
381 ... print(ve)
382 This re-iterable cannot be unpacked.
383 """
384 return source._unpack() if isinstance(source, __Reiterator) else source
387def merge_sorted_and_return_unique(
388 *seqs: Iterable[T]) -> Generator[T, None, None]:
389 """
390 Merge sorted sequences of integers and return only unique values.
392 You can provide multiple sequences, all of which must be sorted.
393 This function then merges them into a single sorted sequence which
394 contains each elemenet at most once.
395 A typical use case would be to combine the result of
396 :func:`pycommons.math.primes.primes` with some pre-defined values into
397 a sorted sequence.
399 Notice that the elements of the sequence must support the less-than
400 operator, i.e., have a `__lt__` dunder method. Otherwise this function
401 will crash.
403 The returned sequence is guaranteed to provide strictly increasing values.
405 :param seqs: the sequences, i.e., some instances of :class:`Iterable` or
406 :class:`Iterator`
407 :returns: a merged sequence of integers
408 :raises TypeError: if any of the provided iterators or any of their
409 elements is `None`, or if any of the elements in `seqs`is not an
410 :class:`Iterable`.
412 >>> list(merge_sorted_and_return_unique([1, 2, 3,], [2, 2]))
413 [1, 2, 3]
415 >>> from pycommons.math.primes import primes
416 >>> list(merge_sorted_and_return_unique(primes(14), [1, 10]))
417 [1, 2, 3, 5, 7, 10, 11, 13]
419 >>> list(merge_sorted_and_return_unique(
420 ... primes(14), primes(17), [1, 2, 10, 100]))
421 [1, 2, 3, 5, 7, 10, 11, 13, 17, 100]
423 >>> try:
424 ... for _ in merge_sorted_and_return_unique(1):
425 ... pass
426 ... except TypeError as te:
427 ... print(te)
428 'int' object is not iterable
430 >>> try:
431 ... for j in merge_sorted_and_return_unique([3], 1):
432 ... print(j)
433 ... except TypeError as te:
434 ... print(te)
435 'int' object is not iterable
437 >>> try:
438 ... for j in merge_sorted_and_return_unique([None], [None]):
439 ... print(j)
440 ... except TypeError as te:
441 ... print(te)
442 Element must not be None.
444 >>> try:
445 ... for j in merge_sorted_and_return_unique([None], [1]):
446 ... print(j)
447 ... except TypeError as te:
448 ... print(te)
449 '<' not supported between instances of 'NoneType' and 'int'
451 >>> try:
452 ... for j in merge_sorted_and_return_unique(None, [1]):
453 ... print(j)
454 ... except TypeError as te:
455 ... print(te)
456 'NoneType' object is not iterable
458 >>> try:
459 ... for j in merge_sorted_and_return_unique([print, len], [repr]):
460 ... print(j)
461 ... except TypeError as te:
462 ... print(te)
463 '<' not supported between instances of 'builtin_function_or_method' \
464and 'builtin_function_or_method'
465 """
466 last: T | None = None
467 for item in merge(*seqs): # type: ignore
468 if item is None:
469 raise TypeError("Element must not be None.")
470 if (last is None) or (last < item): # type: ignore # noqa
471 yield item
472 last = item