Skip to content

Support pickling of figures with aligned x/y labels. #25332

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 28, 2023
Merged

Support pickling of figures with aligned x/y labels. #25332

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions lib/matplotlib/cbook.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -788,6 +788,19 @@ class Grouper:
def __init__(self, init=()):
self._mapping = {weakref.ref(x): [weakref.ref(x)] for x in init}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven't looked at the Grouper much in detail, but would it be possible to use a weakkey dictionary here? Or a weakset for the values?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't check in detail, but weakref.py doesn't seem to make any specific allowance for pickling, so probably something like this PR would remain needed. I guess you could open a separate issue to track the usefulness of switching the implementation to weakkeydicts and weaksets, though.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree you still need to go back and forth between weak and strong during pickling, but I'm a bit confused why you're doing it in Figure's getstate, rather than adding getstate/setstate to the Grouper class itself?

Here was a quick attempt at hacking it in with the WeakKeyDictionary, but now that I think about it, we might actually need the weakset as well here (I overlooked that there was a cycle from key: list(key) in the key-value pairs)
main...greglucas:matplotlib:grouper-pickle

But I think that shows you could probably move your current implementation into getstate/setstate and it would be a little easier to follow all in one place (and allow other Grouper's to take advantage of it).

Copy link
Contributor Author

@anntzer anntzer Feb 27, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, yes. I think I was slightly confused because I also wrote the pickling of shared axes (via _AxesBase._shared_axes) and that one needs specific support (both because the Grouper is (effectively) a global living at the class level and thus doesn't get picked up by pickling, and also because even if it did you wouldn't want to pickle the entire Grouper (and thus all currently existing shared axes) but only the connected components of the axes currently being pickled). But here _align_label_groups is per-figure, so I guess just adding pickling support directly to Grouper is OK (and doesn't harm the above case.

Do you want to take over the issue and push your patch instead? :) [to be clear, the PR here is still correct AFAICT, but @greglucas' solution is more general]

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It will probably be a bit before I could get around to adding tests and thinking more about the proper weak-reference counting and refactoring that area properly, so I say we push forward here for now and leave that as a follow-up PR. Are you willing to move your weak/strong transitions over into Grouper's getstate/setstate instead of the classmethod and figure updates? I think this is good to go then.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, done.


def __getstate__(self):
return {
**vars(self),
# Convert weak refs to strong ones.
"_mapping": {k(): [v() for v in vs] for k, vs in self._mapping.items()},
}

def __setstate__(self, state):
vars(self).update(state)
# Convert strong refs to weak ones.
self._mapping = {weakref.ref(k): [*map(weakref.ref, vs)]
for k, vs in self._mapping.items()}

def __contains__(self, item):
return weakref.ref(item) in self._mapping

Expand Down
4 changes: 4 additions & 0 deletions lib/matplotlib/tests/test_pickle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ def _generate_complete_test_figure(fig_ref):
# Ensure lists also pickle correctly.
plt.subplot(3, 3, 1)
plt.plot(list(range(10)))
plt.ylabel("hello")

plt.subplot(3, 3, 2)
plt.contourf(data, hatches=['//', 'ooo'])
Expand All @@ -68,6 +69,7 @@ def _generate_complete_test_figure(fig_ref):

plt.subplot(3, 3, 4)
plt.imshow(data)
plt.ylabel("hello\nworld!")

plt.subplot(3, 3, 5)
plt.pcolor(data)
Expand All @@ -89,6 +91,8 @@ def _generate_complete_test_figure(fig_ref):
plt.subplot(3, 3, 9)
plt.errorbar(x, x * -0.5, xerr=0.2, yerr=0.4)

fig_ref.align_ylabels() # Test handling of _align_label_groups Groupers.


@mpl.style.context("default")
@check_figures_equal(extensions=["png"])
Expand Down