Merge pull request #5452 from u55/cycler_generic_iterator

WeatherGod · WeatherGod · commit 595868acabd8 · 2015-11-10T13:05:19.000-05:00
Fix axes.set_prop_cycle to handle any generic iterable sequence.
diff --git a/lib/matplotlib/rcsetup.py b/lib/matplotlib/rcsetup.py
@@ -22,6 +22,11 @@
 import operator
 import os
 import warnings
+try:
+    import collections.abc as abc
+except ImportError:
+    # python 2
+    import collections as abc
 from matplotlib.fontconfig_pattern import parse_fontconfig_pattern
 from matplotlib.colors import is_color_like
 
@@ -78,15 +83,19 @@ def f(s):
                     return [scalar_validator(v.strip()) for v in s if v.strip()]
                 else:
                     raise
-        elif type(s) in (list, tuple):
+        # We should allow any generic sequence type, including generators,
+        # Numpy ndarrays, and pandas data structures.  However, unordered
+        # sequences, such as sets, should be allowed but discouraged unless the
+        # user desires pseudorandom behavior.
+        elif isinstance(s, abc.Iterable) and not isinstance(s, abc.Mapping):
             # The condition on this list comprehension will preserve the
             # behavior of filtering out any empty strings (behavior was
             # from the original validate_stringlist()), while allowing
             # any non-string/text scalar values such as numbers and arrays.
             return [scalar_validator(v) for v in s
                     if not isinstance(v, six.string_types) or v]
         else:
-            msg = "'s' must be of type [ string | list | tuple ]"
+            msg = "{0!r} must be of type: string or non-dictionary iterable.".format(s)
             raise ValueError(msg)
     f.__doc__ = scalar_validator.__doc__
     return f
diff --git a/lib/matplotlib/tests/test_cycles.py b/lib/matplotlib/tests/test_cycles.py
@@ -1,6 +1,7 @@
-from matplotlib.testing.decorators import image_comparison
+from matplotlib.testing.decorators import image_comparison, cleanup
 import matplotlib.pyplot as plt
 import numpy as np
+from nose.tools import assert_raises
 
 from cycler import cycler
 
@@ -42,6 +43,27 @@ def test_marker_cycle():
     ax.legend(loc='upper left')
 
 
+# Reuse the image from test_marker_cycle()
+@image_comparison(baseline_images=['marker_cycle'], remove_text=True,
+                  extensions=['png'])
+def test_marker_cycle_keywords():
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    # Test keyword arguments, numpy arrays, and generic iterators
+    ax.set_prop_cycle(color=np.array(['r', 'g', 'y']),
+                      marker=iter(['.', '*', 'x']))
+    xs = np.arange(10)
+    ys = 0.25 * xs + 2
+    ax.plot(xs, ys, label='red dot', lw=4, ms=16)
+    ys = 0.45 * xs + 3
+    ax.plot(xs, ys, label='green star', lw=4, ms=16)
+    ys = 0.65 * xs + 4
+    ax.plot(xs, ys, label='yellow x', lw=4, ms=16)
+    ys = 0.85 * xs + 5
+    ax.plot(xs, ys, label='red2 dot', lw=4, ms=16)
+    ax.legend(loc='upper left')
+
+
 @image_comparison(baseline_images=['lineprop_cycle_basic'], remove_text=True,
                   extensions=['png'])
 def test_linestylecycle_basic():
@@ -104,6 +126,44 @@ def test_fillcycle_ignore():
     ax.legend(loc='upper left')
 
 
+@cleanup
+def test_valid_input_forms():
+    fig, ax = plt.subplots()
+    # These should not raise an error.
+    ax.set_prop_cycle(None)
+    ax.set_prop_cycle(cycler('linewidth', [1, 2]))
+    ax.set_prop_cycle('color', 'rgywkbcm')
+    ax.set_prop_cycle('linewidth', (1, 2))
+    ax.set_prop_cycle('linewidth', [1, 2])
+    ax.set_prop_cycle('linewidth', iter([1, 2]))
+    ax.set_prop_cycle('linewidth', np.array([1, 2]))
+    ax.set_prop_cycle('color', np.array([[1, 0, 0],
+                                         [0, 1, 0],
+                                         [0, 0, 1]]))
+    ax.set_prop_cycle(lw=[1, 2], color=['k', 'w'], ls=['-', '--'])
+    ax.set_prop_cycle(lw=np.array([1, 2]),
+                      color=np.array(['k', 'w']),
+                      ls=np.array(['-', '--']))
+    assert True
+
+
+@cleanup
+def test_invalid_input_forms():
+    fig, ax = plt.subplots()
+    with assert_raises((TypeError, ValueError)):
+        ax.set_prop_cycle(1)
+    with assert_raises((TypeError, ValueError)):
+        ax.set_prop_cycle([1, 2])
+    with assert_raises((TypeError, ValueError)):
+        ax.set_prop_cycle('color', 'fish')
+    with assert_raises((TypeError, ValueError)):
+        ax.set_prop_cycle('linewidth', 1)
+    with assert_raises((TypeError, ValueError)):
+        ax.set_prop_cycle('linewidth', {'1': 1, '2': 2})
+    with assert_raises((TypeError, ValueError)):
+        ax.set_prop_cycle(linewidth=1, color='r')
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
diff --git a/lib/matplotlib/tests/test_rcparams.py b/lib/matplotlib/tests/test_rcparams.py
@@ -280,10 +280,15 @@ def test_validators():
                      ('aardvark, ,', ['aardvark']),
                      (['a', 'b'], ['a', 'b']),
                      (('a', 'b'), ['a', 'b']),
-                     ((1, 2), ['1', '2'])),
-            'fail': ((dict(), ValueError),
-                     (1, ValueError),)
-            },
+                     (iter(['a', 'b']), ['a', 'b']),
+                     (np.array(['a', 'b']), ['a', 'b']),
+                     ((1, 2), ['1', '2']),
+                     (np.array([1, 2]), ['1', '2']),
+                    ),
+         'fail': ((dict(), ValueError),
+                  (1, ValueError),
+                 )
+        },
         {'validator': validate_nseq_int(2),
          'success': ((_, [1, 2])
                      for _ in ('1, 2', [1.5, 2.5], [1, 2],
@@ -353,6 +358,8 @@ def test_validators():
                      (['', 'g', 'blue'], ['g', 'blue']),
                      ([np.array([1, 0, 0]), np.array([0, 1, 0])],
                          np.array([[1, 0, 0], [0, 1, 0]])),
+                     (np.array([[1, 0, 0], [0, 1, 0]]),
+                         np.array([[1, 0, 0], [0, 1, 0]])),
                     ),
          'fail': (('fish', ValueError),
                  ),
