colorbar: remove assert statements (#36)

Author: mathause

mplotutils/__init__.py

@@ -1,8 +1,8 @@
 # flake8: noqa
 
-from . import cartopy_utils, colorbar_utils, mpl_utils
+from . import _colorbar, cartopy_utils, mpl_utils
+from ._colorbar import *
 from .cartopy_utils import *
-from .colorbar_utils import *
 from .mpl_utils import *
 from .xrcompat import *
 

mplotutils/colorbar_utils.py mplotutils/_colorbar.pymplotutils/colorbar_utils.py renamed to mplotutils/_colorbar.py

@@ -1,3 +1,5 @@
+import warnings
+
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -145,8 +147,7 @@ def colorbar(
     plt.colorbar
     """
 
-    orientations = ("vertical", "horizontal")
-    if orientation not in orientations:
+    if orientation not in ("vertical", "horizontal"):
         raise ValueError("orientation must be 'vertical' or 'horizontal'")
 
     k = kwargs.keys()
@@ -156,14 +157,15 @@ def colorbar(
 
     # ensure 'ax' does not end up in plt.colorbar(**kwargs)
     if "ax" in k:
+        if ax2 is not None:
+            raise ValueError("Cannot pass `ax`, and `ax2`")
         # assume it is ax2 (it can't be ax1)
         ax2 = kwargs.pop("ax")
 
     f = ax1.get_figure()
 
-    if ax2 is not None:
-        f2 = ax2.get_figure()
-        assert f == f2, "'ax1' and 'ax2' must belong to the same figure"
+    if ax2 is not None and f != ax2.get_figure():
+        raise ValueError("'ax1' and 'ax2' must belong to the same figure")
 
     cbax = _get_cbax(f)
 
@@ -242,7 +244,7 @@ def _resize_colorbar_vert(
     cbax = f.add_axes([0, 0, 0.1, 0.1])
     cbar = plt.colorbar(h, orientation='vertical', cax=cbax)
 
-    func = mpu.colorbar_utils_resize_colorbar_vert(cbax, ax)
+    func = mpu._resize_colorbar_vert(cbax, ax)
     f.canvas.mpl_connect('draw_event', func)
 
     ax.set_global()
@@ -337,7 +339,7 @@ def _resize_colorbar_horz(
     cbax = f.add_axes([0, 0, 0.1, 0.1])
     cbar = plt.colorbar(h, orientation='horizontal', cax=cbax)
 
-    func = mpu.colorbar_utils._resize_colorbar_horz(cbax, ax)
+    func = mpu._resize_colorbar_horz(cbax, ax)
     f.canvas.mpl_connect('draw_event', func)
 
     ax.set_global()
@@ -379,11 +381,11 @@ def inner(event=None):
             full_width = posn2.x0 - posn1.x0 + posn2.width
 
         pad_scaled = pad * posn1.height
-        size_scaled = size * posn1.height
 
         width = full_width - shrink * full_width
 
         if aspect is None:
+            size_scaled = size * posn1.height
             height = size_scaled
         else:
             figure_aspect = np.divide(*f.get_size_inches())
@@ -414,15 +416,16 @@ def _parse_shift_shrink(shift, shrink):
         if shrink is None:
             shrink = shift
 
-    assert (shift >= 0.0) & (shift <= 1.0), "'shift' must be in 0...1"
-    assert (shrink >= 0.0) & (shrink <= 1.0), "'shrink' must be in 0...1"
+    if (shift < 0.0) or (shift > 1.0):
+        raise ValueError("'shift' must be in 0...1")
+
+    if (shrink < 0.0) or (shrink > 1.0):
+        raise ValueError("'shrink' must be in 0...1")
 
     if shift > shrink:
-        msg = (
-            "Warning: 'shift' is larger than 'shrink', colorbar\n"
-            "will extend beyond the axes!"
+        warnings.warn(
+            "'shift' is larger than 'shrink', colorbar will extend beyond the axes"
         )
-        print(msg)
 
     return shift, shrink
 
@@ -433,16 +436,12 @@ def _parse_shift_shrink(shift, shrink):
 def _parse_size_aspect_pad(size, aspect, pad, orientation):
 
     if (size is not None) and (aspect is not None):
-        raise ValueError("you can only pass one of 'aspect' and 'size'")
+        raise ValueError("Can only pass one of 'aspect' and 'size'")
 
     # default is aspect=20
     if (size is None) and (aspect is None):
         aspect = 20
 
-    # we need a large size so it is not limiting for set_aspect
-    if aspect is not None:
-        size = 10
-
     # default mpl setting
     if pad is None:
         pad = 0.05 if orientation == "vertical" else 0.15

mplotutils/tests/test_colorbar_utils.py

@@ -1,8 +1,8 @@
 import matplotlib.pyplot as plt
 import numpy as np
-from pytest import raises
+import pytest
 
-from mplotutils.colorbar_utils import (
+from mplotutils._colorbar import (
     _get_cbax,
     _parse_shift_shrink,
     _parse_size_aspect_pad,
@@ -25,54 +25,72 @@ def test_parse_shift_shrink():
 
     assert _parse_shift_shrink(0.5, 0.5) == (0.5, 0.5)
 
-    with raises(AssertionError):
+    with pytest.raises(ValueError, match="'shift' must be in 0...1"):
         _parse_shift_shrink(-0.1, 0)
 
-    with raises(AssertionError):
+    with pytest.raises(ValueError, match="'shift' must be in 0...1"):
         _parse_shift_shrink(1.1, 0)
 
-    with raises(AssertionError):
+    with pytest.raises(ValueError, match="'shrink' must be in 0...1"):
         _parse_shift_shrink(0, -0.1)
 
-    with raises(AssertionError):
+    with pytest.raises(ValueError, match="'shrink' must be in 0...1"):
         _parse_shift_shrink(0, 1.1)
 
-
-# =============================================================================
+    with pytest.warns(UserWarning, match="'shift' is larger than 'shrink'"):
+        _parse_shift_shrink(0.6, 0.3)
 
 
 def test_parse_size_aspect_pad():
     """
     size, aspect, pad = _parse_size_aspect_pad(size, aspect, pad, 'horizontal')
     """
 
-    res = _parse_size_aspect_pad(0.1, None, 0.1, "horizontal")
-    exp = (0.1, None, 0.1)
-    assert res == exp
+    with pytest.raises(ValueError, match="Can only pass one of 'aspect' and 'size'"):
+        _parse_size_aspect_pad(1, 1, 0.1, "horizontal")
 
-    res = _parse_size_aspect_pad(None, None, 0.1, "horizontal")
-    exp = (10, 20, 0.1)
-    assert res == exp
+    result = _parse_size_aspect_pad(0.1, None, 0.1, "horizontal")
+    assert result == (0.1, None, 0.1)
 
-    res = _parse_size_aspect_pad(None, 20, 0.1, "horizontal")
-    exp = (10, 20, 0.1)
-    assert res == exp
+    result = _parse_size_aspect_pad(None, None, 0.1, "horizontal")
+    assert result == (None, 20, 0.1)
 
-    with raises(ValueError):
-        _parse_size_aspect_pad(1, 1, 0.1, "horizontal")
+    result = _parse_size_aspect_pad(None, 10, 0.1, "horizontal")
+    assert result == (None, 10, 0.1)
+
+    result = _parse_size_aspect_pad(None, 20, 0.1, "horizontal")
+    assert result == (None, 20, 0.1)
 
-    res = _parse_size_aspect_pad(None, None, None, "horizontal")
-    exp = (10, 20, 0.15)
-    assert res == exp
+    result = _parse_size_aspect_pad(None, None, None, "horizontal")
+    assert result == (None, 20, 0.15)
 
-    res = _parse_size_aspect_pad(None, None, None, "vertical")
-    exp = (10, 20, 0.05)
-    assert res == exp
+    result = _parse_size_aspect_pad(None, None, None, "vertical")
+    assert result == (None, 20, 0.05)
 
 
 # =============================================================================
 
 
+def test_colorbar_differnt_figures():
+
+    _, ax1 = plt.subplots()
+    _, ax2 = plt.subplots()
+
+    h = ax1.pcolormesh([[0, 1]])
+
+    with pytest.raises(ValueError, match="must belong to the same figure"):
+        colorbar(h, ax1, ax2)
+
+
+def test_colorbar_ax_and_ax2_error():
+
+    _, (ax1, ax2, ax3) = plt.subplots(3, 1)
+    h = ax1.pcolormesh([[0, 1]])
+
+    with pytest.raises(ValueError, match="Cannot pass `ax`, and `ax2`"):
+        colorbar(h, ax1, ax2, ax=ax3)
+
+
 def _easy_cbar_vert(**kwargs):
 
     f, ax = plt.subplots()
@@ -363,19 +381,15 @@ def test_colorbar():
     f1, ax1 = plt.subplots()
     h = ax1.pcolormesh([[0, 1]])
 
-    with raises(ValueError):
+    with pytest.raises(ValueError):
         colorbar(h, ax1, orientation="wrong")
 
-    with raises(ValueError):
+    with pytest.raises(ValueError):
         colorbar(h, ax1, anchor=5)
 
-    with raises(ValueError):
+    with pytest.raises(ValueError):
         colorbar(h, ax1, panchor=5)
 
-    with raises(AssertionError):
-        f2, ax2 = plt.subplots()
-        colorbar(h, ax1, ax2)
-
 
 # =============================================================================