Merge pull request #11917 from efiring/contour_level_kwarg_scalar

jklymak · web-flow · commit 7b5ff33daa95 · 2018-08-22T19:41:45.000-07:00
BUG: make arg 'N' and kwarg 'levels' behave the same when scalar
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -1225,21 +1225,15 @@ def _contour_level_args(self, z, args):
         self._auto = False
         if self.levels is None:
             if len(args) == 0:
-                lev = self._autolev(7)
+                levels_arg = 7  # Default, hard-wired.
             else:
-                level_arg = args[0]
-                try:
-                    if isinstance(level_arg, Integral):
-                        lev = self._autolev(level_arg)
-                    else:
-                        lev = np.asarray(level_arg).astype(np.float64)
-                except:
-                    raise TypeError(
-                        "Last {0} arg must give levels; see help({0})"
-                        .format(fn))
-            self.levels = lev
+                levels_arg = args[0]
+        else:
+            levels_arg = self.levels
+        if isinstance(levels_arg, Integral):
+            self.levels = self._autolev(levels_arg)
         else:
-            self.levels = np.asarray(self.levels).astype(np.float64)
+            self.levels = np.asarray(levels_arg).astype(np.float64)
 
         if not self.filled:
             inside = (self.levels > self.zmin) & (self.levels < self.zmax)
diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -130,6 +130,17 @@ def test_contour_empty_levels():
     assert len(record) == 1
 
 
+def test_contour_Nlevels():
+    # A scalar levels arg or kwarg should trigger auto level generation.
+    # https://github.com/matplotlib/matplotlib/issues/11913
+    z = np.arange(12).reshape((3, 4))
+    fig, ax = plt.subplots()
+    cs1 = ax.contour(z, 5)
+    assert len(cs1.levels) > 1
+    cs2 = ax.contour(z, levels=5)
+    assert (cs1.levels == cs2.levels).all()
+
+
 def test_contour_badlevel_fmt():
     # test funny edge case from
     # https://github.com/matplotlib/matplotlib/issues/9742
