ENH : implement __getitem__ for Cycler

tacaswell · tacaswell · commit 126d7df45a5c · 2015-05-05T23:28:16.000-04:00
Implement `__getitem__` for Cycler by broadcasting down to
the lists in the transposed data.

Raises on trying to use anything else but a slice.  Numpy-style
fancy slicing might be worth adding.   Adding integer slicing is
probably of minilmal value.
diff --git a/lib/matplotlib/cycler.py b/lib/matplotlib/cycler.py
@@ -94,6 +94,15 @@ def _from_iter(cls, label, itr):
         ret._keys = set([label])
         return ret
 
+    def __getitem__(self, key):
+        # TODO : maybe add numpy style fancy slicing
+        if isinstance(key, slice):
+            trans = self._transpose()
+            return reduce(add, (cycler(k, v[key])
+                                for k, v in six.iteritems(trans)))
+        else:
+            raise ValueError("Can only use slices with Cycler.__getitem__")
+
     def __iter__(self):
         if self._right is None:
             return iter(self._left)
diff --git a/lib/matplotlib/tests/test_cycler.py b/lib/matplotlib/tests/test_cycler.py
@@ -3,7 +3,7 @@
 
 import six
 from six.moves import zip
-from matplotlib.cycler import cycler
+from matplotlib.cycler import cycler, Cycler
 from nose.tools import assert_equal, assert_raises
 from itertools import product
 from operator import add, iadd, mul, imul
@@ -110,3 +110,18 @@ def test_mul_fails():
     assert_raises(TypeError, mul, c1,  2.0)
     assert_raises(TypeError, mul, c1,  'a')
     assert_raises(TypeError, mul, c1,  [])
+
+
+def test_getitem():
+    c1 = cycler('lw', range(15))
+    for slc in (slice(None, None, None),
+                slice(None, None, -1),
+                slice(1, 5, None),
+                slice(0, 5, 2)):
+        yield _cycles_equal, c1[slc], cycler('lw', range(15)[slc])
+
+
+def test_fail_getime():
+    c1 = cycler('lw', range(15))
+    assert_raises(ValueError, Cycler.__getitem__, c1, 0)
+    assert_raises(ValueError, Cycler.__getitem__, c1, [0, 1])
