Backport of #11917 (#13583)

jklymak · web-flow · commit 6b3c6b997ed6 · 2019-03-04T07:30:09.000-08:00
* Backport from master
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -7,7 +7,9 @@
 import six
 from six.moves import xrange
 
+from numbers import Integral
 import warnings
+
 import matplotlib as mpl
 import numpy as np
 from numpy import ma
@@ -1206,21 +1208,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 type(level_arg) == int:
-                        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
@@ -131,6 +131,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
