8000 Merge pull request #26341 from anntzer/pc · matplotlib/matplotlib@4f65e66 · GitHub
[go: up one dir, main page]

Skip to content

Commit 4f65e66

Browse files
authored
Merge pull request #26341 from anntzer/pc
Fix pickling of axes property cycle.
2 parents eb33b94 + 8e2f62a commit 4f65e66

File tree

4 files changed

+33
-19
lines changed

4 files changed

+33
-19
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/sankey.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -723,7 +723,7 @@ def _get_angle(a, r):
723723
fc = kwargs.pop('fc', kwargs.pop('facecolor', None))
724724
lw = kwargs.pop('lw', kwargs.pop('linewidth', None))
725725
if fc is None:
726-
fc = next(self.ax._get_patches_for_fill.prop_cycler)['color']
726+
fc = self.ax._get_patches_for_fill.get_next_color()
727727
patch = PathPatch(Path(vertices, codes), fc=fc, lw=lw, **kwargs)
728728
self.ax.add_patch(patch)
729729

lib/matplotlib/tests/test_cycles.py

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
import contextlib
2+
from io import StringIO
3+
14
import matplotlib as mpl
25
import matplotlib.pyplot as plt
36
import numpy as np
@@ -120,15 +123,22 @@ def test_valid_input_forms():
120123

121124
def test_cycle_reset():
122125
fig, ax = plt.subplots()
126+
prop0 = StringIO()
127+
prop1 = StringIO()
128+
prop2 = StringIO()
129+
130+
with contextlib.redirect_stdout(prop0):
131+
plt.getp(ax.plot([1, 2], label="label")[0])
123132

124-
# Can't really test a reset because only a cycle object is stored
125-
# but we can test the first item of the cycle.
126-
prop = next(ax._get_lines.prop_cycler)
127133
ax.set_prop_cycle(linewidth=[10, 9, 4])
128-
assert prop != next(ax._get_lines.prop_cycler)
134+
with contextlib.redirect_stdout(prop1):
135+
plt.getp(ax.plot([1, 2], label="label")[0])
136+
assert prop1.getvalue() != prop0.getvalue()
137+
129138
ax.set_prop_cycle(None)
130-
got = next(ax._get_lines.prop_cycler)
131-
assert prop == got
139+
with contextlib.redirect_stdout(prop2):
140+
plt.getp(ax.plot([1, 2], label="label")[0])
141+
assert prop2.getvalue() == prop0.getvalue()
132142

133143

134144
def test_invalid_input_forms():

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