Merge pull request #1139 from dmcdougall/stemarg

mdboom · mdboom · commit 3b649cb0cd68 · 2013-01-04T12:39:19.000-08:00
Make Axes.stem take at least one argument.
diff --git a/CHANGELOG b/CHANGELOG
@@ -21,6 +21,10 @@
            Also added some tests for the normalization class.
            Till Stensitzki
 
+2012-11-12 Make axes.stem take at least one argument.
+           Uses a default range(n) when the first arg not provided.
+           Damon McDougall
+
 2012-11-09 Make plt.subplot() without arguments act as subplot(111) - PI
 
 2012-10-05 Add support for saving animations as animated GIFs. - JVDP
diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py
@@ -5029,20 +5029,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*).
 
@@ -5061,6 +5063,37 @@ def stem(self, x, y, linefmt='b-', markerfmt='bo', basefmt='r-',
             self.cla()
         self.hold(True)
 
+        # Assume there's at least one data array
+        y = np.asarray(args[0], dtype=np.float)
+        args = args[1:]
+
+        # Try a second one
+        try:
+            second = np.asarray(args[0], dtype=np.float)
+            x, y = y, second
+            args = args[1:]
+        except (IndexError, ValueError):
+            # The second array doesn't make sense, or it doesn't exist
+            second = np.arange(len(y))
+            x = second
+
+        # Popping some defaults
+        try:
+            linefmt = kwargs.pop('linefmt', args[0])
+        except IndexError:
+            linefmt = kwargs.pop('linefmt', 'b-')
+        try:
+            markerfmt = kwargs.pop('markerfmt', args[1])
+        except IndexError:
+            markerfmt = kwargs.pop('markerfmt', 'bo')
+        try:
+            basefmt = kwargs.pop('basefmt', args[2])
+        except IndexError:
+            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
@@ -2998,23 +2998,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)
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -924,6 +924,20 @@ def test_hist_stacked_weighted():
     ax = fig.add_subplot(111)
     ax.hist( (d1, d2), weights=(w1,w2), histtype="stepfilled", stacked=True)
 
+@cleanup
+def test_stem_args():
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1)
+
+    x = range(10)
+    y = range(10)
+
+    # Test the call signatures
+    ax.stem(y)
+    ax.stem(x, y)
+    ax.stem(x, y, 'r--')
+    ax.stem(x, y, 'r--', basefmt='b--')
+
 @image_comparison(baseline_images=['transparent_markers'], remove_text=True)
 def test_transparent_markers():
     np.random.seed(0)
