Pickle support. (And Dynamic Inheritance Concept)

Author: isaacrobinson2000

matplotview/_utils.py

@@ -0,0 +1,83 @@
+import functools
+
+def _fix_super_reference(function):
+    """
+    Private utility decorator. Allows a function to be transferred to another
+    class by dynamically updating the local __class__ attribute of the function
+    when called. This allows for use of zero argument super in all methods
+
+    Parameters:
+    -----------
+    function
+        The function to wrap, to allow for the dynamic
+    """
+    @functools.wraps(function)
+    def run_function(self, *args, **kwargs):
+        try:
+            cls_idx = function.__code__.co_freevars.index('__class__')
+            old_value = function.__closure__[cls_idx].cell_contents
+            function.__closure__[cls_idx].cell_contents = type(self)
+            res = function(self, *args, **kwargs)
+            function.__closure__[cls_idx].cell_contents = old_value
+            return res
+        except (AttributeError, ValueError):
+            return function(self, *args, **kwargs)
+
+    return run_function
+
+
+class _WrappingType(type):
+    def __new__(mcs, *args, **kwargs):
+        res = super().__new__(mcs, *args, **kwargs)
+
+        res.__base_wrapping__ = getattr(
+            res, "__base_wrapping__", res.__bases__[0]
+        )
+        res.__instances__ = getattr(res, "__instances__", {})
+
+        return res
+
+    def __getitem__(cls, the_type):
+        if(cls.__instances__ is None):
+            raise TypeError("Already instantiated wrapper!")
+
+        if(the_type == cls.__base_wrapping__):
+            return cls
+
+        if(issubclass(super().__class__, _WrappingType)):
+            return cls._gen_type(super()[the_type])
+
+        if(not issubclass(the_type, cls.__base_wrapping__)):
+            raise TypeError(
+                f"The extension type {the_type} must be a subclass of "
+                f"{cls.__base_wrapping__}"
+            )
+
+        return cls._gen_type(the_type)
+
+    def _gen_type(cls, the_type):
+        if(the_type not in cls.__instances__):
+            cls.__instances__[the_type] = _WrappingType(
+                f"{cls.__name__}[{the_type.__name__}]",
+                (the_type,),
+                {"__instances__": None}
+            )
+            cls._copy_attrs_to(cls.__instances__[the_type])
+
+        return cls.__instances__[the_type]
+
+    def _copy_attrs_to(cls, other):
+        dont_copy = {"__dict__", "__weakref__", "__instances__"}
+
+        for k, v in cls.__dict__.items():
+            if(k not in dont_copy):
+                setattr(
+                    other,
+                    k,
+                    _fix_super_reference(v) if(hasattr(v, "__code__")) else v
+                )
+
+        other.__instances__ = None
+
+    def __iter__(cls):
+        return NotImplemented

matplotview/_view_axes.py

@@ -1,3 +1,4 @@
+import functools
 import itertools
 from typing import Type, List, Optional, Callable, Any
 from matplotlib.axes import Axes
@@ -77,7 +78,12 @@ def do_3d_projection(self) -> float:
 
         return res
 
+def _view_from_pickle(builder, args):
+    res = builder(*args)
+    res.__class__ = view_wrapper(type(res))
+    return res
 
+@functools.lru_cache(None)
 def view_wrapper(axes_class: Type[Axes]) -> Type[Axes]:
     """
     Construct a ViewAxes, which subclasses, or wraps a specific Axes subclass.
@@ -95,15 +101,12 @@ def view_wrapper(axes_class: Type[Axes]) -> Type[Axes]:
         The view axes wrapper for a given axes class, capable of display
         other axes contents...
     """
-
     @docstring.interpd
-    class ViewAxesImpl(axes_class):
+    class View(axes_class):
         """
         An axes which automatically displays elements of another axes. Does not
         require Artists to be plotted twice.
         """
-        __module__ = axes_class.__module__
-
         def __init__(
             self,
             axes_to_view: Axes,
@@ -113,7 +116,7 @@ def __init__(
             filter_function: Optional[Callable[[Artist], bool]] = None,
             **kwargs
         ):
-            f"""
+            """
             Construct a new view axes.
 
             Parameters
@@ -170,10 +173,10 @@ def _init_vars(
         ):
             if(render_depth < 1):
                 raise ValueError(f"Render depth of {render_depth} is invalid.")
-            if(filter_function is None):
-                filter_function = lambda a: True
-            if(not callable(filter_function)):
-                raise ValueError(f"The filter function must be a callable!")
+            if(filter_function is not None and not callable(filter_function)):
+                raise ValueError(
+                    f"The filter function must be a callable or None!"
+                )
 
             self.__view_axes = axes_to_view
             # The current render depth is stored in the figure, so the number
@@ -212,13 +215,14 @@ def get_children(self) -> List[Artist]:
                     for a in itertools.chain(
                         self.__view_axes._children,
                         self.__view_axes.child_axes
-                    ) if(self.__filter_function(a))
+                    ) if(self.__filter_function is None
+                         or self.__filter_function(a))
                 ])
 
                 return init_list
             else:
                 return super().get_children()
-            
+
         def draw(self, renderer: RendererBase = None):
             # It is possible to have two axes which are views of each other
             # therefore we track the number of recursions and stop drawing
@@ -260,6 +264,26 @@ def set_linescaling(self, value: bool):
             """
             self.__scale_lines = value
 
+        def __reduce__(self):
+            builder, args = super().__reduce__()[:2]
+
+            if(type(self) in args):
+                builder = super().__new__
+                args = (type(self).__bases__[0],)
+
+            return (
+                _view_from_pickle,
+                (builder, args),
+                self.__getstate__()
+            )
+
+        def __getstate__(self):
+            state = super().__getstate__()
+            state["__renderer"] = None
+            # We don't support pickling the filter...
+            state["__filter_function"] = None
+            return state
+
         @classmethod
         def from_axes(
             cls,
@@ -276,10 +300,7 @@ def from_axes(
             )
             return axes
 
-    new_name = f"{ViewAxesImpl.__name__}[{axes_class.__name__}]"
-    ViewAxesImpl.__name__ = ViewAxesImpl.__qualname__ = new_name
-
-    return ViewAxesImpl
-
+    View.__name__ = f"{View.__name__}[{axes_class.__name__}]"
+    View.__qualname__ = f"{View.__qualname__}[{axes_class.__name__}]"
 
-ViewAxes = view_wrapper(Axes)
+    return View

matplotview/tests/test_view_obj.py

@@ -0,0 +1,80 @@
+import matplotlib.pyplot as plt
+import pickle
+from matplotview import view, view_wrapper, inset_zoom_axes
+import numpy as np
+
+def to_image(figure):
+    figure.canvas.draw()
+    img = np.frombuffer(figure.canvas.tostring_rgb(), dtype=np.uint8)
+    return img.reshape(figure.canvas.get_width_height()[::-1] + (3,))
+
+
+def test_obj_comparison():
+    from matplotlib.axes import Subplot, Axes
+
+    view_class1 = view_wrapper(Subplot)
+    view_class2 = view_wrapper(Subplot)
+    view_class3 = view_wrapper(Axes)
+
+    assert view_class1 is view_class2
+    assert view_class1 == view_class2
+    assert view_class2 != view_class3
+
+
+def test_subplot_view_pickle():
+    np.random.seed(1)
+    im_data = np.random.rand(30, 30)
+
+    # Test case...
+    fig_test, (ax_test1, ax_test2) = plt.subplots(1, 2)
+
+    ax_test1.plot([i for i in range(10)], "r")
+    ax_test1.add_patch(plt.Circle((3, 3), 1, ec="black", fc="blue"))
+    ax_test1.text(10, 10, "Hello World!", size=14)
+    ax_test1.imshow(im_data, origin="lower", cmap="Blues", alpha=0.5,
+                   interpolation="nearest")
+    ax_test2 = view(ax_test2, ax_test1)
+    ax_test2.set_aspect(ax_test1.get_aspect())
+    ax_test2.set_xlim(ax_test1.get_xlim())
+    ax_test2.set_ylim(ax_test1.get_ylim())
+
+    img_expected = to_image(fig_test)
+
+    saved_fig = pickle.dumps(fig_test)
+    plt.clf()
+
+    fig_test = pickle.loads(saved_fig)
+    img_result = to_image(fig_test)
+
+    assert np.all(img_expected == img_result)
+
+
+def test_zoom_plot_pickle():
+    np.random.seed(1)
+    plt.clf()
+    im_data = np.random.rand(30, 30)
+
+    # Test Case...
+    fig_test = plt.gcf()
+    ax_test = fig_test.gca()
+    ax_test.plot([i for i in range(10)], "r")
+    ax_test.add_patch(plt.Circle((3, 3), 1, ec="black", fc="blue"))
+    ax_test.imshow(im_data, origin="lower", cmap="Blues", alpha=0.5,
+                   interpolation="nearest")
+    axins_test = inset_zoom_axes(ax_test, [0.5, 0.5, 0.48, 0.48])
+    axins_test.set_linescaling(False)
+    axins_test.set_xlim(1, 5)
+    axins_test.set_ylim(1, 5)
+    ax_test.indicate_inset_zoom(axins_test, edgecolor="black")
+
+    fig_test.savefig("before.png")
+    img_expected = to_image(fig_test)
+
+    saved_fig = pickle.dumps(fig_test)
+    plt.clf()
+
+    fig_test = pickle.loads(saved_fig)
+    fig_test.savefig("after.png")
+    img_result = to_image(fig_test)
+
+    assert np.all(img_expected == img_result)

matplotview/tests/test_inset_zoom.py renamed to matplotview/tests/test_view_rendering.py

@@ -136,6 +136,7 @@ def test_3d_view(fig_test, fig_ref):
     ax2_ref.set_ylim(-10, 10)
     ax2_ref.set_zlim(-2, 2)
 
+
 @check_figures_equal()
 def test_polar_view(fig_test, fig_ref):
     r = np.arange(0, 2, 0.01)
@@ -156,6 +157,7 @@ def test_polar_view(fig_test, fig_ref):
     ax_r2.plot(theta, r)
     ax_r2.set_rmax(1)
 
+
 @check_figures_equal()
 def test_map_projection_view(fig_test, fig_ref):
     x = np.linspace(-2.5, 2.5, 20)