__init__.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573
  1. """
  2. Cycler
  3. ======
  4. Cycling through combinations of values, producing dictionaries.
  5. You can add cyclers::
  6. from cycler import cycler
  7. cc = (cycler(color=list('rgb')) +
  8. cycler(linestyle=['-', '--', '-.']))
  9. for d in cc:
  10. print(d)
  11. Results in::
  12. {'color': 'r', 'linestyle': '-'}
  13. {'color': 'g', 'linestyle': '--'}
  14. {'color': 'b', 'linestyle': '-.'}
  15. You can multiply cyclers::
  16. from cycler import cycler
  17. cc = (cycler(color=list('rgb')) *
  18. cycler(linestyle=['-', '--', '-.']))
  19. for d in cc:
  20. print(d)
  21. Results in::
  22. {'color': 'r', 'linestyle': '-'}
  23. {'color': 'r', 'linestyle': '--'}
  24. {'color': 'r', 'linestyle': '-.'}
  25. {'color': 'g', 'linestyle': '-'}
  26. {'color': 'g', 'linestyle': '--'}
  27. {'color': 'g', 'linestyle': '-.'}
  28. {'color': 'b', 'linestyle': '-'}
  29. {'color': 'b', 'linestyle': '--'}
  30. {'color': 'b', 'linestyle': '-.'}
  31. """
  32. from __future__ import annotations
  33. from collections.abc import Hashable, Iterable, Generator
  34. import copy
  35. from functools import reduce
  36. from itertools import product, cycle
  37. from operator import mul, add
  38. # Dict, List, Union required for runtime cast calls
  39. from typing import TypeVar, Generic, Callable, Union, Dict, List, Any, overload, cast
  40. __version__ = "0.12.1"
  41. K = TypeVar("K", bound=Hashable)
  42. L = TypeVar("L", bound=Hashable)
  43. V = TypeVar("V")
  44. U = TypeVar("U")
  45. def _process_keys(
  46. left: Cycler[K, V] | Iterable[dict[K, V]],
  47. right: Cycler[K, V] | Iterable[dict[K, V]] | None,
  48. ) -> set[K]:
  49. """
  50. Helper function to compose cycler keys.
  51. Parameters
  52. ----------
  53. left, right : iterable of dictionaries or None
  54. The cyclers to be composed.
  55. Returns
  56. -------
  57. keys : set
  58. The keys in the composition of the two cyclers.
  59. """
  60. l_peek: dict[K, V] = next(iter(left)) if left != [] else {}
  61. r_peek: dict[K, V] = next(iter(right)) if right is not None else {}
  62. l_key: set[K] = set(l_peek.keys())
  63. r_key: set[K] = set(r_peek.keys())
  64. if l_key & r_key:
  65. raise ValueError("Can not compose overlapping cycles")
  66. return l_key | r_key
  67. def concat(left: Cycler[K, V], right: Cycler[K, U]) -> Cycler[K, V | U]:
  68. r"""
  69. Concatenate `Cycler`\s, as if chained using `itertools.chain`.
  70. The keys must match exactly.
  71. Examples
  72. --------
  73. >>> num = cycler('a', range(3))
  74. >>> let = cycler('a', 'abc')
  75. >>> num.concat(let)
  76. cycler('a', [0, 1, 2, 'a', 'b', 'c'])
  77. Returns
  78. -------
  79. `Cycler`
  80. The concatenated cycler.
  81. """
  82. if left.keys != right.keys:
  83. raise ValueError(
  84. "Keys do not match:\n"
  85. "\tIntersection: {both!r}\n"
  86. "\tDisjoint: {just_one!r}".format(
  87. both=left.keys & right.keys, just_one=left.keys ^ right.keys
  88. )
  89. )
  90. _l = cast(Dict[K, List[Union[V, U]]], left.by_key())
  91. _r = cast(Dict[K, List[Union[V, U]]], right.by_key())
  92. return reduce(add, (_cycler(k, _l[k] + _r[k]) for k in left.keys))
  93. class Cycler(Generic[K, V]):
  94. """
  95. Composable cycles.
  96. This class has compositions methods:
  97. ``+``
  98. for 'inner' products (zip)
  99. ``+=``
  100. in-place ``+``
  101. ``*``
  102. for outer products (`itertools.product`) and integer multiplication
  103. ``*=``
  104. in-place ``*``
  105. and supports basic slicing via ``[]``.
  106. Parameters
  107. ----------
  108. left, right : Cycler or None
  109. The 'left' and 'right' cyclers.
  110. op : func or None
  111. Function which composes the 'left' and 'right' cyclers.
  112. """
  113. def __call__(self):
  114. return cycle(self)
  115. def __init__(
  116. self,
  117. left: Cycler[K, V] | Iterable[dict[K, V]] | None,
  118. right: Cycler[K, V] | None = None,
  119. op: Any = None,
  120. ):
  121. """
  122. Semi-private init.
  123. Do not use this directly, use `cycler` function instead.
  124. """
  125. if isinstance(left, Cycler):
  126. self._left: Cycler[K, V] | list[dict[K, V]] = Cycler(
  127. left._left, left._right, left._op
  128. )
  129. elif left is not None:
  130. # Need to copy the dictionary or else that will be a residual
  131. # mutable that could lead to strange errors
  132. self._left = [copy.copy(v) for v in left]
  133. else:
  134. self._left = []
  135. if isinstance(right, Cycler):
  136. self._right: Cycler[K, V] | None = Cycler(
  137. right._left, right._right, right._op
  138. )
  139. else:
  140. self._right = None
  141. self._keys: set[K] = _process_keys(self._left, self._right)
  142. self._op: Any = op
  143. def __contains__(self, k):
  144. return k in self._keys
  145. @property
  146. def keys(self) -> set[K]:
  147. """The keys this Cycler knows about."""
  148. return set(self._keys)
  149. def change_key(self, old: K, new: K) -> None:
  150. """
  151. Change a key in this cycler to a new name.
  152. Modification is performed in-place.
  153. Does nothing if the old key is the same as the new key.
  154. Raises a ValueError if the new key is already a key.
  155. Raises a KeyError if the old key isn't a key.
  156. """
  157. if old == new:
  158. return
  159. if new in self._keys:
  160. raise ValueError(
  161. f"Can't replace {old} with {new}, {new} is already a key"
  162. )
  163. if old not in self._keys:
  164. raise KeyError(
  165. f"Can't replace {old} with {new}, {old} is not a key"
  166. )
  167. self._keys.remove(old)
  168. self._keys.add(new)
  169. if self._right is not None and old in self._right.keys:
  170. self._right.change_key(old, new)
  171. # self._left should always be non-None
  172. # if self._keys is non-empty.
  173. elif isinstance(self._left, Cycler):
  174. self._left.change_key(old, new)
  175. else:
  176. # It should be completely safe at this point to
  177. # assume that the old key can be found in each
  178. # iteration.
  179. self._left = [{new: entry[old]} for entry in self._left]
  180. @classmethod
  181. def _from_iter(cls, label: K, itr: Iterable[V]) -> Cycler[K, V]:
  182. """
  183. Class method to create 'base' Cycler objects
  184. that do not have a 'right' or 'op' and for which
  185. the 'left' object is not another Cycler.
  186. Parameters
  187. ----------
  188. label : hashable
  189. The property key.
  190. itr : iterable
  191. Finite length iterable of the property values.
  192. Returns
  193. -------
  194. `Cycler`
  195. New 'base' cycler.
  196. """
  197. ret: Cycler[K, V] = cls(None)
  198. ret._left = list({label: v} for v in itr)
  199. ret._keys = {label}
  200. return ret
  201. def __getitem__(self, key: slice) -> Cycler[K, V]:
  202. # TODO : maybe add numpy style fancy slicing
  203. if isinstance(key, slice):
  204. trans = self.by_key()
  205. return reduce(add, (_cycler(k, v[key]) for k, v in trans.items()))
  206. else:
  207. raise ValueError("Can only use slices with Cycler.__getitem__")
  208. def __iter__(self) -> Generator[dict[K, V], None, None]:
  209. if self._right is None:
  210. for left in self._left:
  211. yield dict(left)
  212. else:
  213. if self._op is None:
  214. raise TypeError(
  215. "Operation cannot be None when both left and right are defined"
  216. )
  217. for a, b in self._op(self._left, self._right):
  218. out = {}
  219. out.update(a)
  220. out.update(b)
  221. yield out
  222. def __add__(self, other: Cycler[L, U]) -> Cycler[K | L, V | U]:
  223. """
  224. Pair-wise combine two equal length cyclers (zip).
  225. Parameters
  226. ----------
  227. other : Cycler
  228. """
  229. if len(self) != len(other):
  230. raise ValueError(
  231. f"Can only add equal length cycles, not {len(self)} and {len(other)}"
  232. )
  233. return Cycler(
  234. cast(Cycler[Union[K, L], Union[V, U]], self),
  235. cast(Cycler[Union[K, L], Union[V, U]], other),
  236. zip
  237. )
  238. @overload
  239. def __mul__(self, other: Cycler[L, U]) -> Cycler[K | L, V | U]:
  240. ...
  241. @overload
  242. def __mul__(self, other: int) -> Cycler[K, V]:
  243. ...
  244. def __mul__(self, other):
  245. """
  246. Outer product of two cyclers (`itertools.product`) or integer
  247. multiplication.
  248. Parameters
  249. ----------
  250. other : Cycler or int
  251. """
  252. if isinstance(other, Cycler):
  253. return Cycler(
  254. cast(Cycler[Union[K, L], Union[V, U]], self),
  255. cast(Cycler[Union[K, L], Union[V, U]], other),
  256. product
  257. )
  258. elif isinstance(other, int):
  259. trans = self.by_key()
  260. return reduce(
  261. add, (_cycler(k, v * other) for k, v in trans.items())
  262. )
  263. else:
  264. return NotImplemented
  265. @overload
  266. def __rmul__(self, other: Cycler[L, U]) -> Cycler[K | L, V | U]:
  267. ...
  268. @overload
  269. def __rmul__(self, other: int) -> Cycler[K, V]:
  270. ...
  271. def __rmul__(self, other):
  272. return self * other
  273. def __len__(self) -> int:
  274. op_dict: dict[Callable, Callable[[int, int], int]] = {zip: min, product: mul}
  275. if self._right is None:
  276. return len(self._left)
  277. l_len = len(self._left)
  278. r_len = len(self._right)
  279. return op_dict[self._op](l_len, r_len)
  280. # iadd and imul do not exapand the the type as the returns must be consistent with
  281. # self, thus they flag as inconsistent with add/mul
  282. def __iadd__(self, other: Cycler[K, V]) -> Cycler[K, V]: # type: ignore[misc]
  283. """
  284. In-place pair-wise combine two equal length cyclers (zip).
  285. Parameters
  286. ----------
  287. other : Cycler
  288. """
  289. if not isinstance(other, Cycler):
  290. raise TypeError("Cannot += with a non-Cycler object")
  291. # True shallow copy of self is fine since this is in-place
  292. old_self = copy.copy(self)
  293. self._keys = _process_keys(old_self, other)
  294. self._left = old_self
  295. self._op = zip
  296. self._right = Cycler(other._left, other._right, other._op)
  297. return self
  298. def __imul__(self, other: Cycler[K, V] | int) -> Cycler[K, V]: # type: ignore[misc]
  299. """
  300. In-place outer product of two cyclers (`itertools.product`).
  301. Parameters
  302. ----------
  303. other : Cycler
  304. """
  305. if not isinstance(other, Cycler):
  306. raise TypeError("Cannot *= with a non-Cycler object")
  307. # True shallow copy of self is fine since this is in-place
  308. old_self = copy.copy(self)
  309. self._keys = _process_keys(old_self, other)
  310. self._left = old_self
  311. self._op = product
  312. self._right = Cycler(other._left, other._right, other._op)
  313. return self
  314. def __eq__(self, other: object) -> bool:
  315. if not isinstance(other, Cycler):
  316. return False
  317. if len(self) != len(other):
  318. return False
  319. if self.keys ^ other.keys:
  320. return False
  321. return all(a == b for a, b in zip(self, other))
  322. __hash__ = None # type: ignore
  323. def __repr__(self) -> str:
  324. op_map = {zip: "+", product: "*"}
  325. if self._right is None:
  326. lab = self.keys.pop()
  327. itr = list(v[lab] for v in self)
  328. return f"cycler({lab!r}, {itr!r})"
  329. else:
  330. op = op_map.get(self._op, "?")
  331. msg = "({left!r} {op} {right!r})"
  332. return msg.format(left=self._left, op=op, right=self._right)
  333. def _repr_html_(self) -> str:
  334. # an table showing the value of each key through a full cycle
  335. output = "<table>"
  336. sorted_keys = sorted(self.keys, key=repr)
  337. for key in sorted_keys:
  338. output += f"<th>{key!r}</th>"
  339. for d in iter(self):
  340. output += "<tr>"
  341. for k in sorted_keys:
  342. output += f"<td>{d[k]!r}</td>"
  343. output += "</tr>"
  344. output += "</table>"
  345. return output
  346. def by_key(self) -> dict[K, list[V]]:
  347. """
  348. Values by key.
  349. This returns the transposed values of the cycler. Iterating
  350. over a `Cycler` yields dicts with a single value for each key,
  351. this method returns a `dict` of `list` which are the values
  352. for the given key.
  353. The returned value can be used to create an equivalent `Cycler`
  354. using only `+`.
  355. Returns
  356. -------
  357. transpose : dict
  358. dict of lists of the values for each key.
  359. """
  360. # TODO : sort out if this is a bottle neck, if there is a better way
  361. # and if we care.
  362. keys = self.keys
  363. out: dict[K, list[V]] = {k: list() for k in keys}
  364. for d in self:
  365. for k in keys:
  366. out[k].append(d[k])
  367. return out
  368. # for back compatibility
  369. _transpose = by_key
  370. def simplify(self) -> Cycler[K, V]:
  371. """
  372. Simplify the cycler into a sum (but no products) of cyclers.
  373. Returns
  374. -------
  375. simple : Cycler
  376. """
  377. # TODO: sort out if it is worth the effort to make sure this is
  378. # balanced. Currently it is is
  379. # (((a + b) + c) + d) vs
  380. # ((a + b) + (c + d))
  381. # I would believe that there is some performance implications
  382. trans = self.by_key()
  383. return reduce(add, (_cycler(k, v) for k, v in trans.items()))
  384. concat = concat
  385. @overload
  386. def cycler(arg: Cycler[K, V]) -> Cycler[K, V]:
  387. ...
  388. @overload
  389. def cycler(**kwargs: Iterable[V]) -> Cycler[str, V]:
  390. ...
  391. @overload
  392. def cycler(label: K, itr: Iterable[V]) -> Cycler[K, V]:
  393. ...
  394. def cycler(*args, **kwargs):
  395. """
  396. Create a new `Cycler` object from a single positional argument,
  397. a pair of positional arguments, or the combination of keyword arguments.
  398. cycler(arg)
  399. cycler(label1=itr1[, label2=iter2[, ...]])
  400. cycler(label, itr)
  401. Form 1 simply copies a given `Cycler` object.
  402. Form 2 composes a `Cycler` as an inner product of the
  403. pairs of keyword arguments. In other words, all of the
  404. iterables are cycled simultaneously, as if through zip().
  405. Form 3 creates a `Cycler` from a label and an iterable.
  406. This is useful for when the label cannot be a keyword argument
  407. (e.g., an integer or a name that has a space in it).
  408. Parameters
  409. ----------
  410. arg : Cycler
  411. Copy constructor for Cycler (does a shallow copy of iterables).
  412. label : name
  413. The property key. In the 2-arg form of the function,
  414. the label can be any hashable object. In the keyword argument
  415. form of the function, it must be a valid python identifier.
  416. itr : iterable
  417. Finite length iterable of the property values.
  418. Can be a single-property `Cycler` that would
  419. be like a key change, but as a shallow copy.
  420. Returns
  421. -------
  422. cycler : Cycler
  423. New `Cycler` for the given property
  424. """
  425. if args and kwargs:
  426. raise TypeError(
  427. "cycler() can only accept positional OR keyword arguments -- not both."
  428. )
  429. if len(args) == 1:
  430. if not isinstance(args[0], Cycler):
  431. raise TypeError(
  432. "If only one positional argument given, it must "
  433. "be a Cycler instance."
  434. )
  435. return Cycler(args[0])
  436. elif len(args) == 2:
  437. return _cycler(*args)
  438. elif len(args) > 2:
  439. raise TypeError(
  440. "Only a single Cycler can be accepted as the lone "
  441. "positional argument. Use keyword arguments instead."
  442. )
  443. if kwargs:
  444. return reduce(add, (_cycler(k, v) for k, v in kwargs.items()))
  445. raise TypeError("Must have at least a positional OR keyword arguments")
  446. def _cycler(label: K, itr: Iterable[V]) -> Cycler[K, V]:
  447. """
  448. Create a new `Cycler` object from a property name and iterable of values.
  449. Parameters
  450. ----------
  451. label : hashable
  452. The property key.
  453. itr : iterable
  454. Finite length iterable of the property values.
  455. Returns
  456. -------
  457. cycler : Cycler
  458. New `Cycler` for the given property
  459. """
  460. if isinstance(itr, Cycler):
  461. keys = itr.keys
  462. if len(keys) != 1:
  463. msg = "Can not create Cycler from a multi-property Cycler"
  464. raise ValueError(msg)
  465. lab = keys.pop()
  466. # Doesn't need to be a new list because
  467. # _from_iter() will be creating that new list anyway.
  468. itr = (v[lab] for v in itr)
  469. return Cycler._from_iter(label, itr)