Skip to content

Commit 13a060c

Browse files
chadrikchadrik
committed
stubgen: Fix DocStringParser to capture positional-only and keyword-only argument separators
The DocStringParser was previously discarding `/` and `*` separators when parsing function signatures from docstrings, leading to inaccurate stub generation for functions with positional-only or keyword-only arguments. Changes: - Add `/` and `*` as ArgSig entries in parsed function signatures - Fix state reset logic to properly handle multiple signatures with separators
1 parent 698e910 commit 13a060c

File tree

4 files changed

+160
-40
lines changed

4 files changed

+160
-40
lines changed

‎mypy/stubdoc.py‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -265,14 +265,18 @@ def add_token(self, token: tokenize.TokenInfo) -> None:
265265
self.reset()
266266
return
267267
self.keyword_only = len(self.args)
268+
# Add * as an argument
269+
self.args.append(ArgSig(name="*"))
268270
self.accumulator = ""
269271
else:
270272
if self.accumulator.startswith("*"):
271273
self.keyword_only = len(self.args) + 1
272274
self.arg_name = self.accumulator
273-
if not (
274-
token.string == ")" and self.accumulator.strip() == ""
275-
) and not _ARG_NAME_RE.match(self.arg_name):
275+
if (
276+
not (token.string == ")" and self.accumulator.strip() == "")
277+
and not _ARG_NAME_RE.match(self.arg_name)
278+
and self.arg_name not in ("/", "*")
279+
):
276280
# Invalid argument name.
277281
self.reset()
278282
return
@@ -281,7 +285,7 @@ def add_token(self, token: tokenize.TokenInfo) -> None:
281285
if (
282286
self.state[-1] == STATE_ARGUMENT_LIST
283287
and self.keyword_only is not None
284-
and self.keyword_only == len(self.args)
288+
and self.keyword_only == len(self.args) - 1
285289
and not self.arg_name
286290
):
287291
# Error condition: * must be followed by arguments
@@ -320,8 +324,8 @@ def add_token(self, token: tokenize.TokenInfo) -> None:
320324
self.reset()
321325
return
322326
self.pos_only = len(self.args)
323-
self.state.append(STATE_ARGUMENT_TYPE)
324-
self.accumulator = ""
327+
# Set accumulator to / so it gets processed like a regular argument
328+
self.accumulator = "/"
325329

326330
elif token.type == tokenize.OP and token.string == "->" and self.state[-1] == STATE_INIT:
327331
self.accumulator = ""
@@ -347,6 +351,8 @@ def add_token(self, token: tokenize.TokenInfo) -> None:
347351
self.found = False
348352
self.args = []
349353
self.ret_type = "Any"
354+
self.pos_only = None
355+
self.keyword_only = None
350356
# Leave state as INIT.
351357
else:
352358
self.accumulator += token.string
@@ -356,6 +362,8 @@ def reset(self) -> None:
356362
self.args = []
357363
self.found = False
358364
self.accumulator = ""
365+
self.pos_only = None
366+
self.keyword_only = None
359367

360368
def get_signatures(self) -> list[FunctionSig]:
361369
"""Return sorted copy of the list of signatures found so far."""

‎mypy/test/teststubgen.py‎

Lines changed: 126 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -380,6 +380,40 @@ def test_infer_sig_from_docstring(self) -> None:
380380
],
381381
)
382382

383+
def test_infer_sig_from_docstring_overloads(self) -> None:
384+
385+
assert_equal(
386+
infer_sig_from_docstring("\nfunc(x: int=3) -> int\nfunc(x: str) -> str", "func"),
387+
[
388+
FunctionSig(
389+
name="func", args=[ArgSig(name="x", type="int", default=True)], ret_type="int"
390+
),
391+
FunctionSig(
392+
name="func", args=[ArgSig(name="x", type="str", default=False)], ret_type="str"
393+
),
394+
],
395+
)
396+
397+
assert_equal(
398+
infer_sig_from_docstring("func(x: foo.bar)\nfunc(x: str) -> foo.bar", "func"),
399+
[
400+
FunctionSig(name="func", args=[ArgSig(name="x", type="foo.bar")], ret_type="Any"),
401+
FunctionSig(name="func", args=[ArgSig(name="x", type="str")], ret_type="foo.bar"),
402+
],
403+
)
404+
405+
assert_equal(
406+
infer_sig_from_docstring(
407+
"\nfunc(x: int=3) -> int\nfunc(x: invalid::type<with_template>)", "func"
408+
),
409+
[
410+
FunctionSig(
411+
name="func", args=[ArgSig(name="x", type="int", default=True)], ret_type="int"
412+
),
413+
FunctionSig(name="func", args=[ArgSig(name="x", type=None)], ret_type="Any"),
414+
],
415+
)
416+
383417
def test_infer_sig_from_docstring_duplicate_args(self) -> None:
384418
assert_equal(
385419
infer_sig_from_docstring("\nfunc(x, x) -> str\nfunc(x, y) -> int", "func"),
@@ -435,30 +469,36 @@ def test_infer_sig_from_docstring_args_kwargs_errors(self) -> None:
435469
assert_equal(infer_sig_from_docstring("func(**kwargs, *args) -> int", "func"), [])
436470

437471
def test_infer_sig_from_docstring_positional_only_arguments(self) -> None:
438-
assert_equal(
439-
infer_sig_from_docstring("func(self, /) -> str", "func"),
440-
[FunctionSig(name="func", args=[ArgSig(name="self")], ret_type="str")],
441-
)
442472

443473
assert_equal(
444474
infer_sig_from_docstring("func(self, x, /) -> str", "func"),
445475
[
446476
FunctionSig(
447-
name="func", args=[ArgSig(name="self"), ArgSig(name="x")], ret_type="str"
477+
name="func",
478+
args=[ArgSig(name="self"), ArgSig(name="x"), ArgSig(name="/")],
479+
ret_type="str",
448480
)
449481
],
450482
)
451483

452484
assert_equal(
453485
infer_sig_from_docstring("func(x, /, y) -> int", "func"),
454-
[FunctionSig(name="func", args=[ArgSig(name="x"), ArgSig(name="y")], ret_type="int")],
486+
[
487+
FunctionSig(
488+
name="func",
489+
args=[ArgSig(name="x"), ArgSig(name="/"), ArgSig(name="y")],
490+
ret_type="int",
491+
)
492+
],
455493
)
456494

457495
assert_equal(
458496
infer_sig_from_docstring("func(x, /, *args) -> str", "func"),
459497
[
460498
FunctionSig(
461-
name="func", args=[ArgSig(name="x"), ArgSig(name="*args")], ret_type="str"
499+
name="func",
500+
args=[ArgSig(name="x"), ArgSig(name="/"), ArgSig(name="*args")],
501+
ret_type="str",
462502
)
463503
],
464504
)
@@ -468,51 +508,121 @@ def test_infer_sig_from_docstring_positional_only_arguments(self) -> None:
468508
[
469509
FunctionSig(
470510
name="func",
471-
args=[ArgSig(name="x"), ArgSig(name="kwonly"), ArgSig(name="**kwargs")],
511+
args=[
512+
ArgSig(name="x"),
513+
ArgSig(name="/"),
514+
ArgSig(name="*"),
515+
ArgSig(name="kwonly"),
516+
ArgSig(name="**kwargs"),
517+
],
472518
ret_type="str",
473519
)
474520
],
475521
)
476522

523+
assert_equal(
524+
infer_sig_from_docstring("func(self, /) -> str\nfunc(self, x, /) -> str", "func"),
525+
[
526+
FunctionSig(
527+
name="func", args=[ArgSig(name="self"), ArgSig(name="/")], ret_type="str"
528+
),
529+
FunctionSig(
530+
name="func",
531+
args=[ArgSig(name="self"), ArgSig(name="x"), ArgSig(name="/")],
532+
ret_type="str",
533+
),
534+
],
535+
)
536+
477537
def test_infer_sig_from_docstring_keyword_only_arguments(self) -> None:
478538
assert_equal(
479539
infer_sig_from_docstring("func(*, x) -> str", "func"),
480-
[FunctionSig(name="func", args=[ArgSig(name="x")], ret_type="str")],
540+
[FunctionSig(name="func", args=[ArgSig(name="*"), ArgSig(name="x")], ret_type="str")],
481541
)
482542

483543
assert_equal(
484544
infer_sig_from_docstring("func(x, *, y) -> str", "func"),
485-
[FunctionSig(name="func", args=[ArgSig(name="x"), ArgSig(name="y")], ret_type="str")],
545+
[
546+
FunctionSig(
547+
name="func",
548+
args=[ArgSig(name="x"), ArgSig(name="*"), ArgSig(name="y")],
549+
ret_type="str",
550+
)
551+
],
486552
)
487553

488554
assert_equal(
489555
infer_sig_from_docstring("func(*, x, y) -> str", "func"),
490-
[FunctionSig(name="func", args=[ArgSig(name="x"), ArgSig(name="y")], ret_type="str")],
556+
[
557+
FunctionSig(
558+
name="func",
559+
args=[ArgSig(name="*"), ArgSig(name="x"), ArgSig(name="y")],
560+
ret_type="str",
561+
)
562+
],
491563
)
492564

493565
assert_equal(
494566
infer_sig_from_docstring("func(x, *, kwonly, **kwargs) -> str", "func"),
495567
[
496568
FunctionSig(
497569
name="func",
498-
args=[ArgSig(name="x"), ArgSig(name="kwonly"), ArgSig("**kwargs")],
570+
args=[
571+
ArgSig(name="x"),
572+
ArgSig(name="*"),
573+
ArgSig(name="kwonly"),
574+
ArgSig("**kwargs"),
575+
],
499576
ret_type="str",
500577
)
501578
],
502579
)
503580

581+
assert_equal(
582+
infer_sig_from_docstring(
583+
"func(*, x) -> str\nfunc(x, *, kwonly, **kwargs) -> str", "func"
584+
),
585+
[
586+
FunctionSig(
587+
name="func", args=[ArgSig(name="*"), ArgSig(name="x")], ret_type="str"
588+
),
589+
FunctionSig(
590+
name="func",
591+
args=[
592+
ArgSig(name="x"),
593+
ArgSig(name="*"),
594+
ArgSig(name="kwonly"),
595+
ArgSig("**kwargs"),
596+
],
597+
ret_type="str",
598+
),
599+
],
600+
)
601+
504602
def test_infer_sig_from_docstring_pos_only_and_keyword_only_arguments(self) -> None:
505603
assert_equal(
506604
infer_sig_from_docstring("func(x, /, *, y) -> str", "func"),
507-
[FunctionSig(name="func", args=[ArgSig(name="x"), ArgSig(name="y")], ret_type="str")],
605+
[
606+
FunctionSig(
607+
name="func",
608+
args=[ArgSig(name="x"), ArgSig(name="/"), ArgSig(name="*"), ArgSig(name="y")],
609+
ret_type="str",
610+
)
611+
],
508612
)
509613

510614
assert_equal(
511615
infer_sig_from_docstring("func(x, /, y, *, z) -> str", "func"),
512616
[
513617
FunctionSig(
514618
name="func",
515-
args=[ArgSig(name="x"), ArgSig(name="y"), ArgSig(name="z")],
619+
args=[
620+
ArgSig(name="x"),
621+
ArgSig(name="/"),
622+
ArgSig(name="y"),
623+
ArgSig(name="*"),
624+
ArgSig(name="z"),
625+
],
516626
ret_type="str",
517627
)
518628
],
@@ -525,7 +635,9 @@ def test_infer_sig_from_docstring_pos_only_and_keyword_only_arguments(self) -> N
525635
name="func",
526636
args=[
527637
ArgSig(name="x"),
638+
ArgSig(name="/"),
528639
ArgSig(name="y"),
640+
ArgSig(name="*"),
529641
ArgSig(name="z"),
530642
ArgSig("**kwargs"),
531643
],

‎test-data/pybind11_fixtures/expected_stubs_no_docs/pybind11_fixtures/demo.pyi‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,11 @@ class Point:
1111
degree: ClassVar[Point.AngleUnit] = ...
1212
radian: ClassVar[Point.AngleUnit] = ...
1313
def __init__(self, value: typing.SupportsInt) -> None: ...
14-
def __eq__(self, other: object) -> bool: ...
15-
def __hash__(self) -> int: ...
16-
def __index__(self) -> int: ...
17-
def __int__(self) -> int: ...
18-
def __ne__(self, other: object) -> bool: ...
14+
def __eq__(self, other: object, /) -> bool: ...
15+
def __hash__(self, /) -> int: ...
16+
def __index__(self, /) -> int: ...
17+
def __int__(self, /) -> int: ...
18+
def __ne__(self, other: object, /) -> bool: ...
1919
@property
2020
def name(self) -> str: ...
2121
@property
@@ -28,11 +28,11 @@ class Point:
2828
mm: ClassVar[Point.LengthUnit] = ...
2929
pixel: ClassVar[Point.LengthUnit] = ...
3030
def __init__(self, value: typing.SupportsInt) -> None: ...
31-
def __eq__(self, other: object) -> bool: ...
32-
def __hash__(self) -> int: ...
33-
def __index__(self) -> int: ...
34-
def __int__(self) -> int: ...
35-
def __ne__(self, other: object) -> bool: ...
31+
def __eq__(self, other: object, /) -> bool: ...
32+
def __hash__(self, /) -> int: ...
33+
def __index__(self, /) -> int: ...
34+
def __int__(self, /) -> int: ...
35+
def __ne__(self, other: object, /) -> bool: ...
3636
@property
3737
def name(self) -> str: ...
3838
@property

‎test-data/pybind11_fixtures/expected_stubs_with_docs/pybind11_fixtures/demo.pyi‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,15 @@ class Point:
1717
radian: ClassVar[Point.AngleUnit] = ...
1818
def __init__(self, value: typing.SupportsInt) -> None:
1919
"""__init__(self: pybind11_fixtures.demo.Point.AngleUnit, value: typing.SupportsInt) -> None"""
20-
def __eq__(self, other: object) -> bool:
20+
def __eq__(self, other: object, /) -> bool:
2121
"""__eq__(self: object, other: object, /) -> bool"""
22-
def __hash__(self) -> int:
22+
def __hash__(self, /) -> int:
2323
"""__hash__(self: object, /) -> int"""
24-
def __index__(self) -> int:
24+
def __index__(self, /) -> int:
2525
"""__index__(self: pybind11_fixtures.demo.Point.AngleUnit, /) -> int"""
26-
def __int__(self) -> int:
26+
def __int__(self, /) -> int:
2727
"""__int__(self: pybind11_fixtures.demo.Point.AngleUnit, /) -> int"""
28-
def __ne__(self, other: object) -> bool:
28+
def __ne__(self, other: object, /) -> bool:
2929
"""__ne__(self: object, other: object, /) -> bool"""
3030
@property
3131
def name(self) -> str:
@@ -52,15 +52,15 @@ class Point:
5252
pixel: ClassVar[Point.LengthUnit] = ...
5353
def __init__(self, value: typing.SupportsInt) -> None:
5454
"""__init__(self: pybind11_fixtures.demo.Point.LengthUnit, value: typing.SupportsInt) -> None"""
55-
def __eq__(self, other: object) -> bool:
55+
def __eq__(self, other: object, /) -> bool:
5656
"""__eq__(self: object, other: object, /) -> bool"""
57-
def __hash__(self) -> int:
57+
def __hash__(self, /) -> int:
5858
"""__hash__(self: object, /) -> int"""
59-
def __index__(self) -> int:
59+
def __index__(self, /) -> int:
6060
"""__index__(self: pybind11_fixtures.demo.Point.LengthUnit, /) -> int"""
61-
def __int__(self) -> int:
61+
def __int__(self, /) -> int:
6262
"""__int__(self: pybind11_fixtures.demo.Point.LengthUnit, /) -> int"""
63-
def __ne__(self, other: object) -> bool:
63+
def __ne__(self, other: object, /) -> bool:
6464
"""__ne__(self: object, other: object, /) -> bool"""
6565
@property
6666
def name(self) -> str:

0 commit comments

Comments
 (0)