Make Axes.stem take at least one argument.

dmcdougall · dmcdougall · commit eb953ba8f913 · 2012-11-13T11:37:13.000-06:00
When it takes only one, the abscissae default to np.arange(len(y))
diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py
@@ -5074,20 +5074,22 @@ def broken_barh(self, xranges, yrange, **kwargs):
 
         return col
 
-    def stem(self, x, y, linefmt='b-', markerfmt='bo', basefmt='r-',
-             bottom=None, label=None):
+    def stem(self, *args, **kwargs):
         """
         Create a stem plot.
 
-        Call signature::
+        Call signatures::
 
+          stem(y, linefmt='b-', markerfmt='bo', basefmt='r-')
           stem(x, y, linefmt='b-', markerfmt='bo', basefmt='r-')
 
         A stem plot plots vertical lines (using *linefmt*) at each *x*
         location from the baseline to *y*, and places a marker there
         using *markerfmt*.  A horizontal line at 0 is is plotted using
         *basefmt*.
 
+        If no *x* values are provided, the default is (0, 1, ..., len(y) - 1)
+
         Return value is a tuple (*markerline*, *stemlines*,
         *baseline*).
 
@@ -5104,6 +5106,25 @@ def stem(self, x, y, linefmt='b-', markerfmt='bo', basefmt='r-',
         if not self._hold: self.cla()
         self.hold(True)
 
+        # Assume there's at least one data array
+        y = np.asarray(args[0], dtype=np.float)
+
+        # Try a second one
+        try:
+            second = np.asarray(args[1], dtype=np.float)
+            x, y = y, second
+        except:
+            # The second array doesn't make sense, or it doesn't exist
+            second = np.arange(len(y))
+            x = second
+
+        # Popping some defaults
+        linefmt = kwargs.pop('linefmt', 'b-')
+        markerfmt = kwargs.pop('markerfmt', 'bo')
+        basefmt = kwargs.pop('basefmt', 'r-')
+        bottom = kwargs.pop('bottom', None)
+        label = kwargs.pop('label', None)
+
         markerline, = self.plot(x, y, markerfmt, label="_nolegend_")
 
         if bottom is None:
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -3004,23 +3004,21 @@ def stackplot(x, *args, **kwargs):
         draw_if_interactive()
     finally:
         ax.hold(washold)
-    
+
     return ret
 
 # This function was autogenerated by boilerplate.py.  Do not edit as
 # changes will be lost
 @_autogen_docstring(Axes.stem)
-def stem(x, y, linefmt='b-', markerfmt='bo', basefmt='r-', bottom=None,
-         label=None, hold=None):
+def stem(*args, **kwargs):
     ax = gca()
     # allow callers to override the hold state by passing hold=True|False
     washold = ax.ishold()
-
+    hold = kwargs.pop('hold', None)
     if hold is not None:
         ax.hold(hold)
     try:
-        ret = ax.stem(x, y, linefmt=linefmt, markerfmt=markerfmt,
-                      basefmt=basefmt, bottom=bottom, label=label)
+        ret = ax.stem(*args, **kwargs)
         draw_if_interactive()
     finally:
         ax.hold(washold)
