From 823899a6dbefe23c8d127fc30ad4a96d0ddea146 Mon Sep 17 00:00:00 2001
From: danielballan <daniel.b.allan@gmail.com>
Date: Mon, 15 Feb 2016 23:15:12 -0500
Subject: [PATCH] ENH: Add a cycler-specific chain function.
---
cycler.py | 27 ++++++++++++++++++++++++++-
1 file changed, 26 insertions(+), 1 deletion(-)
diff --git a/cycler.py b/cycler.py
index d69adf1..ea91433 100644
--- a/cycler.py
+++ b/cycler.py
@@ -44,7 +44,7 @@
unicode_literals)
import six
-from itertools import product, cycle
+from itertools import product, cycle, chain as itertools_chain
from six.moves import zip, reduce
from operator import mul, add
import copy
@@ -549,3 +549,28 @@ def _cycler(label, itr):
itr = (v[lab] for v in itr)
return Cycler._from_iter(label, itr)
+
+
+def chain(*args):
+ """
+ Concatenate cyclers.
+
+ This is identical to `itertools.chain`, but it will raise if any arguments
+ do not have a `keys` attribute with the same contents.
+
+ Parameters
+ ----------
+ *args : one or more Cycler objects
+
+ Returns
+ -------
+ cycler : Cycler
+ New `Cycler`, concanetating input `Cycler` objects
+ """
+ _args = iter(args)
+ expected_keys = next(_args.keys)
+ for arg in _args:
+ if arg.keys != expected_keys:
+ raise ValueError("Keys of input cyclers must match. {} does not "
+ "match {}".format(arg.keys, expected_keys))
+ return itertools_chain(*args)