Coverage for python/lsst/utils/iteration.py: 23%
33 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-11-06 20:24 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2022-11-06 20:24 +0000
1# This file is part of utils.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
11#
13from __future__ import annotations
15"""Utilities relating to iterators."""
17__all__ = ["chunk_iterable", "ensure_iterable", "isplit"]
19import itertools
20from collections.abc import Mapping
21from typing import Any, Iterable, Iterator, Tuple, TypeVar
24def chunk_iterable(data: Iterable[Any], chunk_size: int = 1_000) -> Iterator[Tuple[Any, ...]]:
25 """Return smaller chunks of an iterable.
27 Parameters
28 ----------
29 data : iterable of anything
30 The iterable to be chunked. Can be a mapping, in which case
31 the keys are returned in chunks.
32 chunk_size : int, optional
33 The largest chunk to return. Can be smaller and depends on the
34 number of elements in the iterator. Defaults to 1_000.
36 Yields
37 ------
38 chunk : `tuple`
39 The contents of a chunk of the iterator as a `tuple`. A tuple is
40 preferred over an iterator since it is more convenient to tell it is
41 empty and the caller knows it can be sized and indexed.
42 """
43 it = iter(data)
44 while chunk := tuple(itertools.islice(it, chunk_size)):
45 yield chunk
48def ensure_iterable(a: Any) -> Iterable[Any]:
49 """Ensure that the input is iterable.
51 There are multiple cases, when the input is:
53 - iterable, but not a `str` or Mapping -> iterate over elements
54 (e.g. ``[i for i in a]``)
55 - a `str` -> return single element iterable (e.g. ``[a]``)
56 - a Mapping -> return single element iterable
57 - not iterable -> return single element iterable (e.g. ``[a]``).
59 Parameters
60 ----------
61 a : iterable or `str` or not iterable
62 Argument to be converted to an iterable.
64 Returns
65 -------
66 i : `generator`
67 Iterable version of the input value.
68 """
69 if isinstance(a, str):
70 yield a
71 return
72 if isinstance(a, Mapping):
73 yield a
74 return
75 try:
76 yield from a
77 except Exception:
78 yield a
81T = TypeVar("T", str, bytes)
84def isplit(string: T, sep: T) -> Iterator[T]:
85 """Split a string or bytes by separator returning a generator.
87 Parameters
88 ----------
89 string : `str` or `bytes`
90 The string to split into substrings.
91 sep : `str` or `bytes`
92 The separator to use to split the string. Must be the same
93 type as ``string``. Must always be given.
95 Yields
96 ------
97 subset : `str` or `bytes`
98 The next subset extracted from the input until the next separator.
99 """
100 if type(string) is not type(sep):
101 raise TypeError(f"String and separator types must match ({type(string)} != {type(sep)})")
102 begin = 0
103 while True:
104 end = string.find(sep, begin)
105 if end == -1:
106 yield string[begin:]
107 return
108 yield string[begin:end]
109 begin = end + 1