automatically draw on savefig (#98)

* automatically draw on savefig

* update docs

Author: mathause

CHANGELOG.md

@@ -23,6 +23,8 @@
 
 ### Enhancements
 
+- Calling `f.canvas.draw()` before `plt.savefig` is no longer necessary. This now happens
+  automatically ([#98](https://github.com/mathause/mplotutils/pull/98)).
 - Add python 3.12 to list of supported versions ([#89](https://github.com/mathause/mplotutils/pull/89)).
 
 ### Bug fixes

docs/example.py

@@ -49,9 +49,6 @@ def plot_map_mpu():
     # ensure the figure has the correct size
     mpu.set_map_layout(axs)
 
-    # call a draw to ensure the colorbar is correct
-    f.canvas.draw()
-
     f.suptitle("With mplotutils")
 
 

mplotutils/__init__.py

@@ -4,11 +4,14 @@
 
 from . import _colorbar, cartopy_utils, colormaps
 from ._colorbar import *
+from ._savefig import autodraw
 from .cartopy_utils import *
 from .colormaps import *
 from .map_layout import set_map_layout
 from .xrcompat import *
 
+autodraw(True)
+
 try:
     __version__ = _get_version("mplotutils")
 except Exception:

mplotutils/_savefig.py

@@ -0,0 +1,56 @@
+from functools import wraps
+
+from matplotlib.figure import Figure
+
+# ensure the original implementation is not overwritten
+try:
+    savefig_orig
+except NameError:
+    savefig_orig = Figure.savefig
+
+
+def savefig(func):
+
+    @wraps(func)
+    def inner(self, *args, **kwargs):
+
+        # call draw in any case
+        self.canvas.draw()
+
+        return func(self, *args, **kwargs)
+
+    return inner
+
+
+class autodraw:
+
+    def __init__(self, /, toggle):
+
+        self.toggle = toggle
+        if toggle:
+            monkeypatch()
+        else:
+            undo()
+
+    def __enter__(self):
+        return
+
+    def __exit__(self, type, value, traceback):
+
+        if self.toggle:
+            undo()
+        else:
+            monkeypatch()
+
+
+def monkeypatch():
+    # Monkey patch matplotlib to call our savefig instead of the standard
+
+    Figure.savefig = savefig(savefig_orig)
+
+
+def undo():
+    Figure.savefig = savefig_orig
+
+
+# mpu.savefig.autodraw()

mplotutils/tests/test_savefig.py

@@ -0,0 +1,96 @@
+import io
+
+import matplotlib.pyplot as plt
+import pytest
+
+import mplotutils as mpu
+from mplotutils.tests.test_colorbar import create_fig_aspect
+
+from . import figure_context
+
+
+def test_autodraw_orig_func():
+
+    # NOTE plt.Figure.savefig is the overwritten one
+    assert mpu._savefig.savefig_orig is not plt.Figure.savefig
+
+    with mpu.autodraw(False):
+        assert mpu._savefig.savefig_orig is plt.Figure.savefig
+
+    with mpu.autodraw(False):
+        with mpu.autodraw(True):
+            assert mpu._savefig.savefig_orig is not plt.Figure.savefig
+
+    with mpu.autodraw(False):
+        savefig_autodraw = plt.Figure.savefig
+
+    with mpu.autodraw(True):
+        savefig_no_autodraw = plt.Figure.savefig
+
+    assert savefig_no_autodraw is not savefig_autodraw
+
+    mpu.autodraw(True)
+    assert mpu._savefig.savefig_orig is not plt.Figure.savefig
+
+    mpu.autodraw(False)
+    assert mpu._savefig.savefig_orig is plt.Figure.savefig
+
+    # restore to default
+    mpu.autodraw(True)
+
+
+def test_ensure_draw_method_called(monkeypatch):
+    # this is a pseudo-mock test (I think the actual backend would need to be mocked)
+
+    class DrawMethodCalled(Exception):
+        pass
+
+    def draw():
+        raise DrawMethodCalled()
+
+    with figure_context() as f:
+
+        monkeypatch.setattr(f.canvas, "draw", draw)
+
+        # not called when not autodrawing
+        with mpu.autodraw(False):
+            f.savefig(io.BytesIO())
+
+        # called when autodrawing
+        with pytest.raises(DrawMethodCalled):
+            f.savefig(io.BytesIO())
+
+
+def test_saved_figure_not_the_same_vertical():
+
+    with figure_context() as f:
+        create_fig_aspect(aspect=0.5, orientation="vertical")
+
+        file_no_autodraw = io.BytesIO()
+
+        with mpu.autodraw(False):
+            f.savefig(file_no_autodraw)
+
+        file_autodraw = io.BytesIO()
+        with mpu.autodraw(True):
+            f.savefig(file_autodraw)
+
+        assert file_no_autodraw.getvalue() != file_autodraw.getvalue()
+
+
+def test_saved_figure_not_the_same_horizontal():
+
+    with figure_context() as f:
+        create_fig_aspect(aspect=2, orientation="horizontal")
+        # ensure the colorbar is actually on the figure
+        f.subplots_adjust(bottom=0.5)
+
+        file_no_autodraw = io.BytesIO()
+        with mpu.autodraw(False):
+            f.savefig(file_no_autodraw)
+
+        file_autodraw = io.BytesIO()
+        with mpu.autodraw(True):
+            f.savefig(file_autodraw)
+
+        assert file_no_autodraw.getvalue() != file_autodraw.getvalue()