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

33 statements  

« prev     ^ index     » next       coverage.py v6.4.1, created at 2022-06-11 02:29 -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# 

12 

13from __future__ import annotations 

14 

15"""Utilities relating to iterators.""" 

16 

17__all__ = ["chunk_iterable", "ensure_iterable", "isplit"] 

18 

19import itertools 

20from collections.abc import Mapping 

21from typing import Any, Iterable, Iterator, Tuple, TypeVar 

22 

23 

24def chunk_iterable(data: Iterable[Any], chunk_size: int = 1_000) -> Iterator[Tuple[Any, ...]]: 

25 """Return smaller chunks of an iterable. 

26 

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. 

35 

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 

46 

47 

48def ensure_iterable(a: Any) -> Iterable[Any]: 

49 """Ensure that the input is iterable. 

50 

51 There are multiple cases, when the input is: 

52 

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]``). 

58 

59 Parameters 

60 ---------- 

61 a : iterable or `str` or not iterable 

62 Argument to be converted to an iterable. 

63 

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 

79 

80 

81T = TypeVar("T", str, bytes) 

82 

83 

84def isplit(string: T, sep: T) -> Iterator[T]: 

85 """Split a string or bytes by separator returning a generator. 

86 

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. 

94 

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