Coverage for python/lsst/utils/iteration.py : 22%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
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 (
22 Any,
23 Iterable,
24 Iterator,
25 Tuple,
26 TypeVar,
27)
30def chunk_iterable(data: Iterable[Any], chunk_size: int = 1_000) -> Iterator[Tuple[Any, ...]]:
31 """Return smaller chunks of an iterable.
33 Parameters
34 ----------
35 data : iterable of anything
36 The iterable to be chunked. Can be a mapping, in which case
37 the keys are returned in chunks.
38 chunk_size : int, optional
39 The largest chunk to return. Can be smaller and depends on the
40 number of elements in the iterator. Defaults to 1_000.
42 Yields
43 ------
44 chunk : `tuple`
45 The contents of a chunk of the iterator as a `tuple`. A tuple is
46 preferred over an iterator since it is more convenient to tell it is
47 empty and the caller knows it can be sized and indexed.
48 """
49 it = iter(data)
50 while (chunk := tuple(itertools.islice(it, chunk_size))):
51 yield chunk
54def ensure_iterable(a: Any) -> Iterable[Any]:
55 """Ensure that the input is iterable.
57 There are multiple cases, when the input is:
59 - iterable, but not a `str` or Mapping -> iterate over elements
60 (e.g. ``[i for i in a]``)
61 - a `str` -> return single element iterable (e.g. ``[a]``)
62 - a Mapping -> return single element iterable
63 - not iterable -> return single element iterable (e.g. ``[a]``).
65 Parameters
66 ----------
67 a : iterable or `str` or not iterable
68 Argument to be converted to an iterable.
70 Returns
71 -------
72 i : `generator`
73 Iterable version of the input value.
74 """
75 if isinstance(a, str):
76 yield a
77 return
78 if isinstance(a, Mapping):
79 yield a
80 return
81 try:
82 yield from a
83 except Exception:
84 yield a
87T = TypeVar('T', str, bytes)
90def isplit(string: T, sep: T) -> Iterator[T]:
91 """Split a string or bytes by separator returning a generator.
93 Parameters
94 ----------
95 string : `str` or `bytes`
96 The string to split into substrings.
97 sep : `str` or `bytes`
98 The separator to use to split the string. Must be the same
99 type as ``string``. Must always be given.
101 Yields
102 ------
103 subset : `str` or `bytes`
104 The next subset extracted from the input until the next separator.
105 """
106 if type(string) is not type(sep):
107 raise TypeError(f"String and separator types must match ({type(string)} != {type(sep)})")
108 begin = 0
109 while True:
110 end = string.find(sep, begin)
111 if end == -1:
112 yield string[begin:]
113 return
114 yield string[begin:end]
115 begin = end + 1