1 from __future__ import annotations

2

3 from typing import (

4 Iterable,

5 Any,

6 Generator,

7 Callable,

8 Union,

9)

10

11 import typing

12 from typing import (

13 Literal,

14 Mapping,

15)

16

17

18 WhenMissing = Literal["except", "skip", "include"]

19

20

21 def empty_sequence_if_attr_false(

22 itr: Iterable[Any],

23 attr_owner: Any,

24 attr_name: str,

25) -> Iterable[Any]:

26 """Returns `itr` if `attr_owner` has the attribute `attr_name` and it boolean casts to `True`. Returns an empty sequence otherwise.

27

28 Particularly useful for optionally inserting delimiters into a sequence depending on an `TokenizerElement` attribute.

29

30 # Parameters:

31 - `itr: Iterable[Any]`

32 The iterable to return if the attribute is `True`.

33 - `attr_owner: Any`

34 The object to check for the attribute.

35 - `attr_name: str`

36 The name of the attribute to check.

37

38 # Returns:

39 - `itr: Iterable` if `attr_owner` has the attribute `attr_name` and it boolean casts to `True`, otherwise an empty sequence.

40 - `()` an empty sequence if the attribute is `False` or not present.

41 """

42 return itr if bool(getattr(attr_owner, attr_name, False)) else ()

43

44

45 def flatten(it: Iterable[Any], levels_to_flatten: int | None = None) -> Generator:

46 """

47 Flattens an arbitrarily nested iterable.

48 Flattens all iterable data types except for `str` and `bytes`.

49

50 # Returns

51 Generator over the flattened sequence.

52

53 # Parameters

54 - `it`: Any arbitrarily nested iterable.

55 - `levels_to_flatten`: Number of levels to flatten by, starting at the outermost layer. If `None`, performs full flattening.

56 """

57 for x in it:

58 # TODO: swap type check with more general check for __iter__() or __next__() or whatever

59 if (

60 hasattr(x, "__iter__")

61 and not isinstance(x, (str, bytes))

62 and (levels_to_flatten is None or levels_to_flatten > 0)

63 ):

64 yield from flatten(

65 x, None if levels_to_flatten is None else levels_to_flatten - 1

66 )

67 else:

68 yield x

69

70

71# string-like operations on lists

72# --------------------------------------------------------------------------------

73

74

75 def list_split(lst: list, val: Any) -> list[list]:

76 """split a list into sublists by `val`. similar to "a_b_c".split("_")

77

78 ```python

79 >>> list_split([1,2,3,0,4,5,0,6], 0)

80 [[1, 2, 3], [4, 5], [6]]

81 >>> list_split([0,1,2,3], 0)

82 [[], [1, 2, 3]]

83 >>> list_split([1,2,3], 0)

84 [[1, 2, 3]]

85 >>> list_split([], 0)

86 [[]]

87 ```

88

89 """

90

91 if len(lst) == 0:

92 return [[]]

93

94 output: list[list] = [

95 [],

96 ]

97

98 for x in lst:

99 if x == val:

100 output.append([])

101 else:

102 output[-1].append(x)

103 return output

104

105

106 def list_join(lst: list, factory: Callable) -> list:

107 """add a *new* instance of `factory()` between each element of `lst`

108

109 ```python

110 >>> list_join([1,2,3], lambda : 0)

111 [1,0,2,0,3]

112 >>> list_join([1,2,3], lambda: [time.sleep(0.1), time.time()][1])

113 [1, 1600000000.0, 2, 1600000000.1, 3]

114 ```

115 """

116

117 if len(lst) == 0:

118 return []

119

120 output: list = [

121 lst[0],

122 ]

123

124 for x in lst[1:]:

125 output.append(factory())

126 output.append(x)

127

128 return output

129

130

131# applying mappings

132# --------------------------------------------------------------------------------

133

134_AM_K = typing.TypeVar("_AM_K")

135_AM_V = typing.TypeVar("_AM_V")

136

137

138 def apply_mapping(

139 mapping: Mapping[_AM_K, _AM_V],

140 iter: Iterable[_AM_K],

141 when_missing: WhenMissing = "skip",

142) -> list[Union[_AM_K, _AM_V]]:

143 """Given an iterable and a mapping, apply the mapping to the iterable with certain options

144

145 Gotcha: if `when_missing` is invalid, this is totally fine until a missing key is actually encountered.

146

147 Note: you can use this with `muutils.kappa.Kappa` if you want to pass a function instead of a dict

148

149 # Parameters:

150 - `mapping : Mapping[_AM_K, _AM_V]`

151 must have `__contains__` and `__getitem__`, both of which take `_AM_K` and the latter returns `_AM_V`

152 - `iter : Iterable[_AM_K]`

153 the iterable to apply the mapping to

154 - `when_missing : WhenMissing`

155 what to do when a key is missing from the mapping -- this is what distinguishes this function from `map`

156 you can choose from `"skip"`, `"include"` (without converting), and `"except"`

157 (defaults to `"skip"`)

158

159 # Returns:

160 return type is one of:

161 - `list[_AM_V]` if `when_missing` is `"skip"` or `"except"`

162 - `list[Union[_AM_K, _AM_V]]` if `when_missing` is `"include"`

163

164 # Raises:

165 - `KeyError` : if the item is missing from the mapping and `when_missing` is `"except"`

166 - `ValueError` : if `when_missing` is invalid

167 """

168 output: list[Union[_AM_K, _AM_V]] = list()

169 item: _AM_K

170 for item in iter:

171 if item in mapping:

172 output.append(mapping[item])

173 continue

174 if when_missing == "skip":

175 continue

176 elif when_missing == "include":

177 output.append(item)

178 elif when_missing == "except":

179 raise KeyError(f"item {item} is missing from mapping {mapping}")

180 else:

181 raise ValueError(

182 f"invalid value for {when_missing = }\n{item = }\n{mapping = }"

183 )

184 return output

185

186

187 def apply_mapping_chain(

188 mapping: Mapping[_AM_K, Iterable[_AM_V]],

189 iter: Iterable[_AM_K],

190 when_missing: WhenMissing = "skip",

191) -> list[Union[_AM_K, _AM_V]]:

192 """Given an iterable and a mapping, chain the mappings together

193

194 Gotcha: if `when_missing` is invalid, this is totally fine until a missing key is actually encountered.

195

196 Note: you can use this with `muutils.kappa.Kappa` if you want to pass a function instead of a dict

197

198 # Parameters:

199 - `mapping : Mapping[_AM_K, Iterable[_AM_V]]`

200 must have `__contains__` and `__getitem__`, both of which take `_AM_K` and the latter returns `Iterable[_AM_V]`

201 - `iter : Iterable[_AM_K]`

202 the iterable to apply the mapping to

203 - `when_missing : WhenMissing`

204 what to do when a key is missing from the mapping -- this is what distinguishes this function from `map`

205 you can choose from `"skip"`, `"include"` (without converting), and `"except"`

206 (defaults to `"skip"`)

207

208 # Returns:

209 return type is one of:

210 - `list[_AM_V]` if `when_missing` is `"skip"` or `"except"`

211 - `list[Union[_AM_K, _AM_V]]` if `when_missing` is `"include"`

212

213 # Raises:

214 - `KeyError` : if the item is missing from the mapping and `when_missing` is `"except"`

215 - `ValueError` : if `when_missing` is invalid

216

217 """

218 output: list[Union[_AM_K, _AM_V]] = list()

219 item: _AM_K

220 for item in iter:

221 if item in mapping:

222 output.extend(mapping[item])

223 continue

224 if when_missing == "skip":

225 continue

226 elif when_missing == "include":

227 output.append(item)

228 elif when_missing == "except":

229 raise KeyError(f"item {item} is missing from mapping {mapping}")

230 else:

231 raise ValueError(

232 f"invalid value for {when_missing = }\n{item = }\n{mapping = }"

233 )

234 return output