Fix pickling of axes property cycle.

anntzer · anntzer · commit 8e2f62a7790e · 2023-07-19T11:04:26.000+02:00
Instead of relying on itertools.cycle (which is unpicklable, and leads
to the axes property cycle being reset to the rcparams value upon
unpickling), use a plain, trivially picklable, list-and-index
representation (the property cycle is guaranteed to be finite because
Axes.set_prop_cycle passes its argument through matplotlib.cycler, which
already iterates through all cycler values for validation purposes).

Unfortunately this is not enough to really get rid of the
get_next_color/get_next_color_func awkwardness, which also serves the
additional purpose of not unnecessarily advancing the prop_cycle when it
does not contain a "color" key and a color is requested.
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -2,7 +2,6 @@
 from contextlib import ExitStack
 import functools
 import inspect
-import itertools
 import logging
 from numbers import Real
 from operator import attrgetter
@@ -224,18 +223,11 @@ def __init__(self, command='plot'):
         self.command = command
         self.set_prop_cycle(None)
 
-    def __getstate__(self):
-        # note: it is not possible to pickle a generator (and thus a cycler).
-        return {'command': self.command}
-
-    def __setstate__(self, state):
-        self.__dict__ = state.copy()
-        self.set_prop_cycle(None)
-
     def set_prop_cycle(self, cycler):
         if cycler is None:
             cycler = mpl.rcParams['axes.prop_cycle']
-        self.prop_cycler = itertools.cycle(cycler)
+        self._idx = 0
+        self._cycler_items = [*cycler]
         self._prop_keys = cycler.keys  # This should make a copy
 
     def __call__(self, axes, *args, data=None, **kwargs):
@@ -315,7 +307,9 @@ def get_next_color(self):
         """Return the next color in the cycle."""
         if 'color' not in self._prop_keys:
             return 'k'
-        return next(self.prop_cycler)['color']
+        c = self._cycler_items[self._idx]['color']
+        self._idx = (self._idx + 1) % len(self._cycler_items)
+        return c
 
     def _getdefaults(self, ignore, kw):
         """
@@ -328,7 +322,8 @@ def _getdefaults(self, ignore, kw):
         if any(kw.get(k, None) is None for k in prop_keys):
             # Need to copy this dictionary or else the next time around
             # in the cycle, the dictionary could be missing entries.
-            default_dict = next(self.prop_cycler).copy()
+            default_dict = self._cycler_items[self._idx].copy()
+            self._idx = (self._idx + 1) % len(self._cycler_items)
             for p in ignore:
                 default_dict.pop(p, None)
         else:
diff --git a/lib/matplotlib/sankey.py b/lib/matplotlib/sankey.py
@@ -723,7 +723,7 @@ def _get_angle(a, r):
             fc = kwargs.pop('fc', kwargs.pop('facecolor', None))
             lw = kwargs.pop('lw', kwargs.pop('linewidth', None))
         if fc is None:
-            fc = next(self.ax._get_patches_for_fill.prop_cycler)['color']
+            fc = self.ax._get_patches_for_fill.get_next_color()
         patch = PathPatch(Path(vertices, codes), fc=fc, lw=lw, **kwargs)
         self.ax.add_patch(patch)
 
diff --git a/lib/matplotlib/tests/test_cycles.py b/lib/matplotlib/tests/test_cycles.py
@@ -1,3 +1,6 @@
+import contextlib
+from io import StringIO
+
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
@@ -120,15 +123,22 @@ def test_valid_input_forms():
 
 def test_cycle_reset():
     fig, ax = plt.subplots()
+    prop0 = StringIO()
+    prop1 = StringIO()
+    prop2 = StringIO()
+
+    with contextlib.redirect_stdout(prop0):
+        plt.getp(ax.plot([1, 2], label="label")[0])
 
-    # Can't really test a reset because only a cycle object is stored
-    # but we can test the first item of the cycle.
-    prop = next(ax._get_lines.prop_cycler)
     ax.set_prop_cycle(linewidth=[10, 9, 4])
-    assert prop != next(ax._get_lines.prop_cycler)
+    with contextlib.redirect_stdout(prop1):
+        plt.getp(ax.plot([1, 2], label="label")[0])
+    assert prop1.getvalue() != prop0.getvalue()
+
     ax.set_prop_cycle(None)
-    got = next(ax._get_lines.prop_cycler)
-    assert prop == got
+    with contextlib.redirect_stdout(prop2):
+        plt.getp(ax.plot([1, 2], label="label")[0])
+    assert prop2.getvalue() == prop0.getvalue()
 
 
 def test_invalid_input_forms():
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -292,3 +292,12 @@ def test_dynamic_norm():
 def test_vertexselector():
     line, = plt.plot([0, 1], picker=True)
     pickle.loads(pickle.dumps(VertexSelector(line)))
+
+
+def test_cycler():
+    ax = plt.figure().add_subplot()
+    ax.set_prop_cycle(c=["c", "m", "y", "k"])
+    ax.plot([1, 2])
+    ax = pickle.loads(pickle.dumps(ax))
+    l, = ax.plot([3, 4])
+    assert l.get_color() == "m"
