ensure correct gca (#128)

Author: mathause

mplotutils/mpl.py

@@ -13,3 +13,18 @@ def _get_renderer(fig):
     raise AttributeError(
         f"Could not find a renderer for the '{backend}' backend. Please raise an issue"
     )
+
+
+def _maybe_gca(**kwargs):
+
+    import matplotlib.pyplot as plt
+
+    # can call gcf unconditionally: either it exists or would be created by plt.axes
+    f = plt.gcf()
+
+    # only call gca if an active axes exists
+    if f.axes:
+        # can not pass kwargs to active axes
+        return plt.gca()
+
+    return plt.axes(**kwargs)

mplotutils/tests/test_maybe_gca.py

@@ -0,0 +1,30 @@
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+
+from mplotutils.mpl import _maybe_gca
+
+from . import figure_context
+
+
+def test_maybe_gca():
+
+    with figure_context():
+        ax = _maybe_gca(aspect=1)
+
+        assert isinstance(ax, mpl.axes.Axes)
+        assert ax.get_aspect() == 1
+
+    with figure_context():
+        ax = _maybe_gca(aspect=1)
+
+        assert isinstance(ax, mpl.axes.Axes)
+        assert ax.get_aspect() == 1
+
+    with figure_context():
+        existing_axes = plt.axes()
+        ax = _maybe_gca(aspect=1)
+
+        # re-uses the existing axes
+        assert existing_axes == ax
+        # kwargs are ignored when reusing axes
+        assert ax.get_aspect() == "auto"