Coverage for python/lsst/utils/iteration.py: 8%
105 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:16 -0700
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#
13"""Utilities relating to iterators."""
15from __future__ import annotations
17__all__ = ["chunk_iterable", "ensure_iterable", "isplit", "sequence_to_string"]
19import itertools
20from collections.abc import Iterable, Iterator, Mapping
21from typing import Any, TypeGuard, 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 : `~collections.abc.Iterable`
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
112def _extract_numeric_suffix(s: str) -> tuple[str, int | None]:
113 """Extract the numeric suffix from a string.
115 Returns the prefix and the numeric suffix as an integer, if present.
117 For example:
118 'node1' -> ('node', 1)
119 'node' -> ('node', None)
120 'node123abc' -> ('node123abc', None)
122 Parameters
123 ----------
124 s : str
125 The string to extract the numeric suffix from.
127 Returns
128 -------
129 suffix : str
130 The numeric suffix of the string, if any.
131 """
132 index = len(s)
133 while index > 0 and s[index - 1].isdigit():
134 index -= 1
135 prefix = s[:index]
136 suffix = s[index:]
137 if suffix:
138 return prefix, int(suffix)
139 else:
140 return s, None
143def _add_pair_to_name(val_name: list[str], val0: int | str, val1: int | str, stride: int = 1) -> None:
144 """Format a pair of values (val0 and val1) and appends the result to
145 val_name.
147 This helper function takes the starting and ending values of a sequence
148 and formats them into a compact string representation, considering the
149 stride and whether the values are integers or strings with common
150 prefixes.
152 Parameters
153 ----------
154 val_name : List[str]
155 The list to append the formatted string to.
156 val0 : [int, str]
157 The starting value of the sequence.
158 val1 : [int, str]
159 The ending value of the sequence.
160 stride : int, optional
161 The stride or difference between consecutive numbers in the
162 sequence. Defaults to 1.
163 """
164 if isinstance(val0, str) and isinstance(val1, str):
165 prefix0, num_suffix0 = _extract_numeric_suffix(val0)
166 prefix1, num_suffix1 = _extract_numeric_suffix(val1)
167 if prefix0 == prefix1 and num_suffix0 is not None and num_suffix1 is not None:
168 if num_suffix0 == num_suffix1:
169 dvn = val0
170 else:
171 dvn = f"{val0}..{val1}"
172 if stride > 1:
173 dvn += f":{stride}"
174 else:
175 dvn = val0 if val0 == val1 else f"{val0}^{val1}"
176 else:
177 sval0 = str(val0)
178 sval1 = str(val1)
179 if val0 == val1:
180 dvn = sval0
181 elif isinstance(val0, int) and isinstance(val1, int):
182 if val1 == val0 + stride:
183 dvn = f"{sval0}^{sval1}"
184 else:
185 dvn = f"{sval0}..{sval1}"
186 if stride > 1:
187 dvn += f":{stride}"
188 else:
189 dvn = f"{sval0}^{sval1}"
190 val_name.append(dvn)
193def _is_list_of_ints(values: list[int | str]) -> TypeGuard[list[int]]:
194 """Check if a list is composed entirely of integers.
196 Parameters
197 ----------
198 values : List[int, str]:
199 The list of values to check.
201 Returns
202 -------
203 is_ints : bool
204 True if all values are integers, False otherwise.
205 """
206 return all(isinstance(v, int) for v in values)
209def sequence_to_string(values: list[int | str]) -> str:
210 """Convert a list of integers or strings into a compact string
211 representation by merging consecutive values or sequences.
213 This function takes a list of integers or strings, sorts them, identifies
214 sequences where consecutive numbers differ by a consistent stride, or
215 strings with common prefixes, and returns a string that compactly
216 represents these sequences. Consecutive numbers are merged into ranges, and
217 strings with common prefixes are handled to produce a concise
218 representation.
220 >>> getNameOfSet([1, 2, 3, 5, 7, 8, 9])
221 '1..3^5^7..9'
222 >>> getNameOfSet(["node1", "node2", "node3"])
223 'node1..node3'
224 >>> getNameOfSet([10, 20, 30, 40])
225 '10..40:10'
227 Parameters
228 ----------
229 values : list[int, str]:
230 A list of items to be compacted. Must all be of the same type.
232 Returns
233 -------
234 sequence_as_string : str
235 A compact string representation of the input list.
237 Notes
238 -----
239 - The function handles both integers and strings.
240 - For strings with common prefixes, only the differing suffixes are
241 considered.
242 - The stride is determined as the minimum difference between
243 consecutive numbers.
244 - Strings without common prefixes are listed individually.
245 """
246 if not values:
247 return ""
249 values = sorted(set(values))
251 pure_ints_or_pure_strings = all(isinstance(item, int) for item in values) or all(
252 isinstance(item, str) for item in values
253 )
254 if not pure_ints_or_pure_strings:
255 types = {type(item) for item in values}
256 raise TypeError(f"All items in the input list must be either integers or strings, got {types}")
258 # Determine the stride for integers
259 stride = 1
260 if len(values) > 1 and _is_list_of_ints(values):
261 differences = [values[i + 1] - values[i] for i in range(len(values) - 1)]
262 stride = min(differences) if differences else 1
263 stride = max(stride, 1)
265 val_name: list[str] = []
266 val0 = values[0]
267 val1 = val0
268 for val in values[1:]:
269 if isinstance(val, int):
270 assert isinstance(val1, int)
271 if val == val1 + stride:
272 val1 = val
273 else:
274 _add_pair_to_name(val_name, val0, val1, stride)
275 val0 = val
276 val1 = val0
277 elif isinstance(val, str):
278 assert isinstance(val1, str)
279 prefix1, num_suffix1 = _extract_numeric_suffix(val1)
280 prefix, num_suffix = _extract_numeric_suffix(val)
281 if prefix1 == prefix and num_suffix1 is not None and num_suffix is not None:
282 if num_suffix == num_suffix1 + stride:
283 val1 = val
284 else:
285 _add_pair_to_name(val_name, val0, val1)
286 val0 = val
287 val1 = val0
288 else:
289 _add_pair_to_name(val_name, val0, val1)
290 val0 = val
291 val1 = val0
293 _add_pair_to_name(val_name, val0, val1, stride)
295 return "^".join(val_name)