You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

558 lines
16 KiB

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