Improve the tests, docstrings, argument checking

efiring · efiring · commit 49902e963ddf · 2020-05-31T22:32:12.000-10:00
The substantive difference is that a ValueError is now raised
if BoundaryNorm is called with ncolors &lt; the total number of
bins, including any extensions.
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -530,7 +530,7 @@ def __call__(self, X, alpha=None, bytes=False):
         """
         Parameters
         ----------
-        X : float, ndarray
+        X : float or int, ndarray or scalar
             The data value(s) to convert to RGBA.
             For floats, X should be in the interval ``[0.0, 1.0]`` to
             return the RGBA values ``X*100`` percent along the Colormap line.
@@ -1410,7 +1410,7 @@ class BoundaryNorm(Normalize):
     interpolation, but using integers seems simpler, and reduces the number of
     conversions back and forth between integer and floating point.
     """
-    def __init__(self, boundaries, ncolors, clip=False, extend='neither'):
+    def __init__(self, boundaries, ncolors, clip=False, *, extend='neither'):
         """
         Parameters
         ----------
@@ -1436,14 +1436,18 @@ def __init__(self, boundaries, ncolors, clip=False, extend='neither'):
             `~matplotlib.colorbar.Colorbar` will be drawn with
             the triangle extension on the left or lower end.
 
+        Returns
+        -------
+        int16 scalar or array
+
         Notes
         -----
         *boundaries* defines the edges of bins, and data falling within a bin
         is mapped to the color with the same index.
 
-        If the number of bins, including any extensions, doesn't equal
-        *ncolors*, the color is chosen by linear interpolation of the
-        bin number onto color numbers.
+        If the number of bins, including any extensions, is less than
+        *ncolors*, the color index is chosen by linear interpolation, mapping
+        the ``[0, nbins - 1]`` range onto the ``[0, ncolors - 1]`` range.
         """
         if clip and extend != 'neither':
             raise ValueError("'clip=True' is not compatible with 'extend'")
@@ -1462,6 +1466,11 @@ def __init__(self, boundaries, ncolors, clip=False, extend='neither'):
             self._offset = 1
         if extend in ('max', 'both'):
             self._N += 1
+        if self._N > self.Ncmap:
+            raise ValueError(f"There are {self._N} color bins including "
+                             f"extensions, but ncolors = {ncolors}; "
+                             "ncolors must equal or exceed the number of "
+                             "bins")
 
     def __call__(self, value, clip=None):
         if clip is None:
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -265,85 +265,83 @@ def test_BoundaryNorm():
     vals = np.ma.masked_invalid([np.Inf])
     assert np.all(bn(vals).mask)
 
-    # Testing extend keyword
-    bounds = [1, 2, 3]
-    cmap = cm.get_cmap('jet')
+    # Incompatible extend and clip
+    with pytest.raises(ValueError, match="not compatible"):
+        mcolors.BoundaryNorm(np.arange(4), 5, extend='both', clip=True)
 
-    refnorm = mcolors.BoundaryNorm([0] + bounds + [4], cmap.N)
+    # Too small ncolors argument
+    with pytest.raises(ValueError, match="ncolors must equal or exceed"):
+        mcolors.BoundaryNorm(np.arange(4), 2)
+
+    with pytest.raises(ValueError, match="ncolors must equal or exceed"):
+        mcolors.BoundaryNorm(np.arange(4), 3, extend='min')
+
+    with pytest.raises(ValueError, match="ncolors must equal or exceed"):
+        mcolors.BoundaryNorm(np.arange(4), 4, extend='both')
+
+    # Testing extend keyword, with interpolation (large cmap)
+    bounds = [1, 2, 3]
+    cmap = cm.get_cmap('viridis')
     mynorm = mcolors.BoundaryNorm(bounds, cmap.N, extend='both')
-    x = np.random.random(100) + 1.5
-    np.testing.assert_array_equal(refnorm(x), mynorm(x))
+    refnorm = mcolors.BoundaryNorm([0] + bounds + [4], cmap.N)
+    x = np.random.randn(100) * 10 + 2
+    ref = refnorm(x)
+    ref[ref == 0] = -1
+    ref[ref == cmap.N - 1] = cmap.N
+    assert_array_equal(mynorm(x), ref)
 
-    # Min and max
+    # Without interpolation
     cmref = mcolors.ListedColormap(['blue', 'red'])
     cmref.set_over('black')
     cmref.set_under('white')
-
     cmshould = mcolors.ListedColormap(['white', 'blue', 'red', 'black'])
-    cmshould.set_over(cmshould(cmshould.N))
-    cmshould.set_under(cmshould(0))
 
     refnorm = mcolors.BoundaryNorm(bounds, cmref.N)
     mynorm = mcolors.BoundaryNorm(bounds, cmshould.N, extend='both')
-    np.testing.assert_array_equal(refnorm.vmin, mynorm.vmin)
-    np.testing.assert_array_equal(refnorm.vmax, mynorm.vmax)
-    x = [-1, 1.2, 2.3, 9.6]
-    np.testing.assert_array_equal(cmshould([0, 1, 2, 3]), cmshould(mynorm(x)))
-    x = np.random.randn(100) * 10 + 2
-    np.testing.assert_array_equal(cmref(refnorm(x)), cmshould(mynorm(x)))
+    assert mynorm.vmin == refnorm.vmin
+    assert mynorm.vmax == refnorm.vmax
 
-    np.testing.assert_array_equal(-1, mynorm(-1))
-    np.testing.assert_array_equal(1, mynorm(1.1))
-    np.testing.assert_array_equal(4, mynorm(12))
+    assert mynorm(bounds[0] - 0.1) == -1  # under
+    assert mynorm(bounds[0] + 0.1) == 1   # first bin -> second color
+    assert mynorm(bounds[-1] - 0.1) == cmshould.N - 2  # next-to-last color
+    assert mynorm(bounds[-1] + 0.1) == cmshould.N  # over
 
-    # Test raises
-    with pytest.raises(ValueError):
-        mcolors.BoundaryNorm(bounds, cmref.N, extend='both', clip=True)
+    x = [-1, 1.2, 2.3, 9.6]
+    assert_array_equal(cmshould(mynorm(x)), cmshould([0, 1, 2, 3]))
+    x = np.random.randn(100) * 10 + 2
+    assert_array_equal(cmshould(mynorm(x)), cmref(refnorm(x)))
 
     # Just min
     cmref = mcolors.ListedColormap(['blue', 'red'])
     cmref.set_under('white')
     cmshould = mcolors.ListedColormap(['white', 'blue', 'red'])
-    cmshould.set_under(cmshould(0))
 
-    np.testing.assert_array_equal(2, cmref.N)
-    np.testing.assert_array_equal(3, cmshould.N)
+    assert cmref.N == 2
+    assert cmshould.N == 3
     refnorm = mcolors.BoundaryNorm(bounds, cmref.N)
     mynorm = mcolors.BoundaryNorm(bounds, cmshould.N, extend='min')
-    np.testing.assert_array_equal(refnorm.vmin, mynorm.vmin)
-    np.testing.assert_array_equal(refnorm.vmax, mynorm.vmax)
+    assert mynorm.vmin == refnorm.vmin
+    assert mynorm.vmax == refnorm.vmax
     x = [-1, 1.2, 2.3]
-    np.testing.assert_array_equal(cmshould([0, 1, 2]), cmshould(mynorm(x)))
+    assert_array_equal(cmshould(mynorm(x)), cmshould([0, 1, 2]))
     x = np.random.randn(100) * 10 + 2
-    np.testing.assert_array_equal(cmref(refnorm(x)), cmshould(mynorm(x)))
+    assert_array_equal(cmshould(mynorm(x)), cmref(refnorm(x)))
 
     # Just max
     cmref = mcolors.ListedColormap(['blue', 'red'])
     cmref.set_over('black')
     cmshould = mcolors.ListedColormap(['blue', 'red', 'black'])
-    cmshould.set_over(cmshould(2))
 
-    np.testing.assert_array_equal(2, cmref.N)
-    np.testing.assert_array_equal(3, cmshould.N)
+    assert cmref.N == 2
+    assert cmshould.N == 3
     refnorm = mcolors.BoundaryNorm(bounds, cmref.N)
     mynorm = mcolors.BoundaryNorm(bounds, cmshould.N, extend='max')
-    np.testing.assert_array_equal(refnorm.vmin, mynorm.vmin)
-    np.testing.assert_array_equal(refnorm.vmax, mynorm.vmax)
+    assert mynorm.vmin == refnorm.vmin
+    assert mynorm.vmax == refnorm.vmax
     x = [1.2, 2.3, 4]
-    np.testing.assert_array_equal(cmshould([0, 1, 2]), cmshould(mynorm(x)))
+    assert_array_equal(cmshould(mynorm(x)), cmshould([0, 1, 2]))
     x = np.random.randn(100) * 10 + 2
-    np.testing.assert_array_equal(cmref(refnorm(x)), cmshould(mynorm(x)))
-
-    # General case
-    bounds = [1, 2, 3, 4]
-    cmap = cm.get_cmap('jet')
-    mynorm = mcolors.BoundaryNorm(bounds, cmap.N, extend='both')
-    refnorm = mcolors.BoundaryNorm([-100] + bounds + [100], cmap.N)
-    x = np.random.randn(100) * 10 - 5
-    ref = refnorm(x)
-    ref = np.where(ref == 0, -1, ref)
-    ref = np.where(ref == cmap.N-1, cmap.N, ref)
-    np.testing.assert_array_equal(ref, mynorm(x))
+    assert_array_equal(cmshould(mynorm(x)), cmref(refnorm(x)))
 
 
 @pytest.mark.parametrize("vmin,vmax", [[-1, 2], [3, 1]])
