Fix passing stem markerfmt positionally when locs are not given

timhoffm · timhoffm · commit a65f7e6f5b32 · 2022-06-09T23:25:46.000+02:00
the signature is stem([locs], heads, linefmt=None, ...)

So we should support both:

```
stem(heads, linefmt='r')
stem(heads, 'r')
```

We had a kwonly deprecation for 3.5 that was aiming at
`stem([locs], heads, *, linefmt=None, ...)`
but now I'd rather relax this to
`stem([locs], heads, linefmt=None, *, ...)`
because it's reasonable to still support `stem(heads, 'r')`. That's
analogous to `plot(y, 'r')`. The code overhead for supporting both
positional and keyword passing for the single `linefmt` parameter
is bearable.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -2900,12 +2900,15 @@ def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None, bottom=0,
             heads, = args
             locs = np.arange(len(heads))
             args = ()
+        elif isinstance(args[1], str):
+            heads, *args = args
+            locs = np.arange(len(heads))
         else:
             locs, heads, *args = args
-        if args:
+        if len(args) > 1:
             _api.warn_deprecated(
                 "3.5",
-                message="Passing the linefmt parameter positionally is "
+                message="Passing the markerfmt parameter positionally is "
                         "deprecated since Matplotlib %(since)s; the "
                         "parameter will become keyword-only %(removal)s.")
 
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3801,16 +3801,23 @@ def test_stem(use_line_collection):
 
 
 def test_stem_args():
+    def _assert_equal(stem_container, expected):
+        x, y = map(list, stem_container.markerline.get_data())
+        assert x == expected[0]
+        assert y == expected[1]
+
     fig, ax = plt.subplots()
 
-    x = list(range(10))
-    y = list(range(10))
+    x = [1, 3, 5]
+    y = [9, 8, 7]
 
     # Test the call signatures
-    ax.stem(y)
-    ax.stem(x, y)
-    ax.stem(x, y, linefmt='r--')
-    ax.stem(x, y, linefmt='r--', basefmt='b--')
+    _assert_equal(ax.stem(y), expected=([0, 1, 2], y))
+    _assert_equal(ax.stem(x, y), expected=(x, y))
+    _assert_equal(ax.stem(x, y, linefmt='r--'), expected=(x, y))
+    _assert_equal(ax.stem(x, y, linefmt='r--', basefmt='b--'), expected=(x, y))
+    _assert_equal(ax.stem(y, linefmt='r--'), expected=([0, 1, 2], y))
+    _assert_equal(ax.stem(y, 'r--'), expected=([0, 1, 2], y))
 
 
 def test_stem_dates():
