Vector colorbar compatibility

CHANGELOG.rst

@@ -9,6 +9,8 @@ The format is based on `Keep a Changelog <https://keepachangelog.com/en/1.0.0/>`
 Unreleased
 ----------
 
+-  Modified ep.colorbar() to work with GeoDataFrame plots (@nkorinek, #557)
+
 0.9.2
 -----
 

earthpy/plot.py

@@ -16,7 +16,7 @@
 import earthpy.spatial as es
 
 
-def colorbar(mapobj, size="3%", pad=0.09):
+def colorbar(mapobj, size="3%", pad=0.09, return_cax=False):
     """Adjust colorbar height to match the matplotlib axis height.
 
     NOTE: This function requires matplotlib v 3.0.1 or greater or v 2.9 or
@@ -31,6 +31,9 @@ def colorbar(mapobj, size="3%", pad=0.09):
         The percent width of the colorbar relative to the plot.
     pad : int (default = 0.09)
         The space between the plot and the color bar.
+    return_cax : bool (default = False)
+        Boolean to return the cax created before it is applied to the axis.
+        Can be used to help adjust colorbars on non-array plots.
 
     Returns
     -------
@@ -70,7 +73,10 @@ def colorbar(mapobj, size="3%", pad=0.09):
     fig = ax.figure
     divider = make_axes_locatable(ax)
     cax = divider.append_axes("right", size=size, pad=pad)
-    return fig.colorbar(mapobj, cax=cax)
+    if return_cax:
+        return cax
+    else:
+        return fig.colorbar(mapobj, cax=cax)
 
 
 def _plot_image(

earthpy/tests/test_plot.py

@@ -21,3 +21,12 @@ def test_colorbar_raises_value_error():
     with pytest.raises(AttributeError, match="requires a matplotlib"):
         ep.colorbar(list())
     plt.close()
+
+
+def test_colorbar_returns_cax(basic_geometry_gdf):
+    """Test that colorbar works with geodataframes when cax is returned."""
+    f, ax = plt.subplots(figsize=(5, 5))
+    cb_ax = ep.colorbar(ax, return_cax=True)
+    im = basic_geometry_gdf.plot(ax=ax, legend=True, cax=cb_ax)
+    assert ax.get_position().height == im.axes.get_position().height
+    plt.close(f)