colorbar: ensure same gca (#136)

* colorbar: ensure same gca

* changelog and comment

Author: mathause

CHANGELOG.md

@@ -30,6 +30,10 @@
 - `set_map_layout` now raises an explicit error when the figure contains SubFigure ([#121](https://github.com/mathause/mplotutils/pull/121)).
 - Test upstream dependencies and fix compatibility with the upcoming pandas v3  ([#133](https://github.com/mathause/mplotutils/pull/133)).
 
+### Bug fixes
+
+- Ensure the current axes (`plt.gca()`) is not changed by calling `mpu.colorbar(...)` ([#136](https://github.com/mathause/mplotutils/pull/136)).
+
 ## v0.5.0 (27.03.2024)
 
 Version v0.5.0 aligns passing multiple axes to `colorbar` with matplotlib.

mplotutils/_colorbar.py

@@ -211,7 +211,10 @@ def colorbar(
     if not all(f == ax.get_figure() for ax in axs):
         raise TypeError("All passed axes must belong to the same figure")
 
+    gca = plt.gca()
     cbax = _get_cbax(f)
+    # ensure mpu.colorbar does not change the current axes
+    plt.sca(gca)
 
     cbar = plt.colorbar(mappable, orientation=orientation, cax=cbax, **kwargs)
 

mplotutils/tests/test_colorbar.py

@@ -235,6 +235,23 @@ def test_colorbar_2d_array():
         assert_position(cbar, expected)
 
 
+def test_gca_is_same():
+
+    with figure_context():
+        h, ax = create_figure_subplots(1, 1)
+        mpu.colorbar(h, ax, size=0.2, pad=0)
+
+        gca = plt.gca()
+        assert gca is ax
+
+    with figure_context():
+        h, axs = create_figure_subplots(1, 2)
+        mpu.colorbar(h, axs, size=0.2, pad=0, orientation="horizontal")
+
+        gca = plt.gca()
+        assert gca is axs[1]
+
+
 def test_colorbar_vertical_aspect():
     with figure_context(figsize=(5, 5)):
         # test pad=0, aspect=5