[Bug]: make_norm_from_scale should create picklable classes even when used in-line.

Bug summary

The new matplotlib.colors.make_norm_from_scale helper dynamically generates a norm class from a scale class. Currently, in the codebase, it is only used as a decorator to create "toplevel" classes (e.g., it is used to generate LogNorm from LogScale, etc.), but it can also be used within other functions to dynamically generate a norm class based on a user-given arbitrary scale (see #20752 for an example of application, which is however not necessary to understand for what follows). In the latter case, the dynamically generate class is currently not picklable (because pickling of classes relies on the existence of global names, see e.g. "classes are pickled by named reference"). It would be generally useful to get rid of this restriction, which can be done by implementing the __reduce__ protocol; there's already other examples in the codebase of dynamically generated classes that use the same mechanism).

I'm tagging this as "good first issue" because there's no API design and the eng goal is clear, but medium (perhaps hard) difficulty because it requires somewhat sophisticated understanding of the details of the pickling process.

Code for reproduction

pickle.dumps(matplotlib.colors.make_norm_from_scale(matplotlib.scale.LogitScale, matplotlib.colors.Normalize))

Actual outcome

Can't pickle <class 'matplotlib.colors.Normalize'>: it's not the same object as matplotlib.colors.Normalize

(Note the additional confusion here: there's two classes that are both named matplotlib.colors.Normalize -- the original one and the dynamically generated one -- but they are different.)

Expected outcome

A correct round-trippable pickle.

Operating system

No response

Matplotlib Version

master (unreleased, pre 3.5)

Matplotlib Backend

No response

Python version

No response

Jupyter version

No response

Other libraries

No response

Installation

source

Conda channel

No response

Hey, I want to try and fix this bug. Will you assign me or should I just start working on it?

Go for it.

I think this wasn't fully solved for the built-in classes of matplotlib.colors. Adding this test to test_dynamic_norm:

    # Also test the builtin classes
    norm = mpl.colors.LogNorm(vmin=0.01, vmax=100)

    assert type(pickle.loads(pickle.dumps(norm))) \
        == type(norm)

Fails with:

E       _pickle.PicklingError: Can't pickle <class 'matplotlib.colors.LogNorm'>: it's not the same object as matplotlib.colors.LogNorm

I tried to look into this a bit, and I'm not sure you can use a decorator in this situation to wrap a class definition and keep it picklable at the module-level. I thought you could possibly update the definition in the globals() dictionary, but even this doesn't seem to help here. Explicitly calling the decorator as a function, makes things picklable, but then you lose the extra class docstring information that was defined. i.e. here I get information about functools.partial in the docstring.

LogNorm = make_norm_from_scale(functools.partial(scale.LogScale, nonpositive="mask"))

I guess I'm somewhat curious if there is some more magic here that can be done to make this actually work... but I'm also wondering how much this is actually gaining over just explicitly defining these named classes in the module (and defining self._scale / self._trf manually). Note that I'm not saying to remove the function for easily making these norms from scales, but rather just explicitly defining the 4 module-level Norms that use the decorator may be the better route here.

I guess the way out may(?) be to just explicitly check for the condition we want:

diff --git i/lib/matplotlib/colors.py w/lib/matplotlib/colors.py
index 6d126e6725..95e63a7459 100644
--- i/lib/matplotlib/colors.py
+++ w/lib/matplotlib/colors.py
@@ -1530,6 +1530,10 @@ def _make_norm_from_scale(scale_cls, base_norm_cls, bound_init_signature):
 
     class Norm(base_norm_cls):
         def __reduce__(self):
+            import importlib
+            if type(self) is getattr(importlib.import_module(type(self).__module__),
+                                     type(self).__qualname__):
+                return (_create_empty_object_of_class, (type(self),), self.__dict__)
             return (_picklable_norm_constructor,
                     (scale_cls, base_norm_cls, bound_init_signature),
                     self.__dict__)
@@ -1610,6 +1614,10 @@ def _picklable_norm_constructor(*args):
     return cls.__new__(cls)
 
 
+def _create_empty_object_of_class(cls):
+    return cls.__new__(cls)
+
+
 @make_norm_from_scale(
     scale.FuncScale,
     init=lambda functions, vmin=None, vmax=None, clip=False: None)

(then _picklable_norm_constructor can be rewritten to use _create_empty_object_of_class) (also, probably needs some error checking e.g. getattr failing should just fall back to the old path)

At least the following now holds:

from matplotlib.colors import LogNorm; from pickle import *; print(type(loads(dumps(LogNorm()))) is LogNorm)

As usual, feel free to pick up the patch.

Navigation Menu

Repository navigation

[Bug]: make_norm_from_scale should create picklable classes even when used in-line. #20755

Bug summary

Code for reproduction

Actual outcome

Expected outcome

Operating system

Matplotlib Version

Matplotlib Backend

Python version

Jupyter version

Other libraries

Installation

Conda channel

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[Bug]: make_norm_from_scale should create picklable classes even when used in-line. #20755

Description

Bug summary

Code for reproduction

Actual outcome

Expected outcome

Operating system

Matplotlib Version

Matplotlib Backend

Python version

Jupyter version

Other libraries

Installation

Conda channel

Activity

Bart-del commented on Sep 7, 2021

anntzer commented on Sep 7, 2021

greglucas commented on Feb 27, 2022

greglucas commented on Mar 20, 2022

anntzer commented on Mar 20, 2022

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions