8000 Fix pickling of axes property cycle. · matplotlib/matplotlib@18d4a4b · GitHub
[go: up one dir, main page]
Skip to content

Commit 18d4a4b

Browse files
anntzeranntzer
committed
Fix pickling of axes property cycle.
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.
1 parent 2a4d905 commit 18d4a4b

File tree

2 files changed

+16
-12
lines changed

2 files changed

+16
-12
lines changed

‎lib/matplotlib/axes/_base.py

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
from contextlib import ExitStack
33
import functools
44
import inspect
5-
import itertools
65
import logging
76
from numbers import Real
87
from operator import attrgetter
@@ -224,18 +223,11 @@ def __init__(self, command='plot'):
224223
self.command = command
225224
self.set_prop_cycle(None)
226225

227-
def __getstate__(self):
228-
# note: it is not possible to pickle a generator (and thus a cycler).
229-
return {'command': self.command}
230-
231-
def __setstate__(self, state):
232-
self.__dict__ = state.copy()
233-
self.set_prop_cycle(None)
234-
235226
def set_prop_cycle(self, cycler):
236227
if cycler is None:
237228
cycler = mpl.rcParams['axes.prop_cycle']
238-
self.prop_cycler = itertools.cycle(cycler)
229+
self._idx = 0
230+
self._cycler_items = [*cycler]
239231
self._prop_keys = cycler.keys # This should make a copy
240232

241233
def __call__(self, axes, *args, data=None, **kwargs):
@@ -315,7 +307,9 @@ def get_next_color(self):
315307
"""Return the next color in the cycle."""
316308
if 'color' not in self._prop_keys:
317309
return 'k'
318-
return next(self.prop_cycler)['color']
310+
c = self._cycler_items[self._idx]['color']
311+
self._idx = (self._idx + 1) % len(self._cycler_items)
312+
return c
319313

320314
def _getdefaults(self, ignore, kw):
321315
"""
@@ -328,7 +322,8 @@ def _getdefaults(self, ignore, kw):
328322
if any(kw.get(k, None) is None for k in prop_keys):
329323
# Need to copy this dictionary or else the next time around
330324
# in the cycle, the dictionary could be missing entries.
331-
default_dict = next(self.prop_cycler).copy()
325+
default_dict = self._cycler_items[self._idx].copy()
326+
self._idx = (self._idx + 1) % len(self._cycler_items)
332327
for p in ignore:
333328
default_dict.pop(p, None)
334329
else:

‎lib/matplotlib/tests/test_pickle.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -292,3 +292,12 @@ def test_dynamic_norm():
292292
def test_vertexselector():
293293
line, = plt.plot([0, 1], picker=True)
294294
pickle.loads(pickle.dumps(VertexSelector(line)))
295+
296+
297+
def test_cycler():
298+
ax = plt.figure().add_subplot()
299+
ax.set_prop_cycle(c=["c", "m", "y", "k"])
300+
ax.plot([1, 2])
301+
ax = pickle.loads(pickle.dumps(ax))
302+
l, = ax.plot([3, 4])
303+
assert l.get_color() == "m"

0 commit comments

Comments
 (0)
0