diff --git a/CHANGELOG b/CHANGELOG index 19f1e06299fc..111a1708f053 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -2,6 +2,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 index ff7187b79740..116c409214fc 100644 --- a/lib/matplotlib/axes.py +++ b/lib/matplotlib/axes.py @@ -5074,13 +5074,13 @@ 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* @@ -5088,6 +5088,8 @@ def stem(self, x, y, linefmt='b-', markerfmt='bo', basefmt='r-', 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,37 @@ 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) + 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 index 4dd2f44504ba..436ea0cc63bd 100644 --- 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) diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py index 1b7b15b31898..550db3101a21 100644 --- a/lib/matplotlib/tests/test_axes.py +++ b/lib/matplotlib/tests/test_axes.py @@ -901,6 +901,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)