fixup! Support positional and keyword-only arguments

Author: pganssle-google

mypy/stubdoc.py

@@ -254,13 +254,21 @@ def add_token(self, token: tokenize.TokenInfo) -> None:
                 self.arg_type = self.accumulator
                 self.state.pop()
             elif self.state[-1] == STATE_ARGUMENT_LIST:
-                self.arg_name = self.accumulator
-                if not (
-                    token.string == ")" and self.accumulator.strip() == ""
-                ) and not _ARG_NAME_RE.match(self.arg_name):
-                    # Invalid argument name.
-                    self.reset()
-                    return
+                if self.accumulator == "*":
+                    if self.keyword_only is not None:
+                        # Error condition: cannot have * twice
+                        self.reset()
+                        return
+                    self.keyword_only = len(self.args)
+                    self.accumulator = ""
+                else:
+                    self.arg_name = self.accumulator
+                    if not (
+                        token.string == ")" and self.accumulator.strip() == ""
+                    ) and not _ARG_NAME_RE.match(self.arg_name):
+                        # Invalid argument name.
+                        self.reset()
+                        return
 
             if token.string == ")":
                 if (
@@ -305,13 +313,10 @@ def add_token(self, token: tokenize.TokenInfo) -> None:
                     self.reset()
                     return
                 self.pos_only = len(self.args)
+                self.state.append(STATE_ARGUMENT_TYPE)
+                self.accumulator = ""
             else:
-                if self.keyword_only is not None:
-                    # * is not allowed after *
-                    self.reset()
-                    return
-                self.keyword_only = len(self.args)
-            self.state.append(STATE_ARGUMENT_TYPE)
+                self.accumulator = "*"
 
         elif token.type == tokenize.OP and token.string == "->" and self.state[-1] == STATE_INIT:
             self.accumulator = ""

mypy/test/teststubgen.py

@@ -399,6 +399,28 @@ def test_infer_sig_from_docstring_bad_indentation(self) -> None:
             None,
         )
 
+    def test_infer_sig_from_docstring_args_kwargs(self) -> None:
+        assert_equal(
+            infer_sig_from_docstring("func(*args, **kwargs) -> int", "func"),
+            [
+                FunctionSig(
+                    name="func",
+                    args=[ArgSig(name="*args"), ArgSig(name="**kwargs")],
+                    ret_type="int",
+                )
+            ],
+        )
+
+        assert_equal(
+            infer_sig_from_docstring("func(*args) -> int", "func"),
+            [FunctionSig(name="func", args=[ArgSig(name="*args")], ret_type="int")],
+        )
+
+        assert_equal(
+            infer_sig_from_docstring("func(**kwargs) -> int", "func"),
+            [FunctionSig(name="func", args=[ArgSig(name="**kwargs")], ret_type="int")],
+        )
+
     def test_infer_sig_from_docstring_positional_only_arguments(self) -> None:
         assert_equal(
             infer_sig_from_docstring("func(self, /) -> str", "func"),