Skip to content

Commit 67cc059

Browse files
authored
stubgen: Support ParamSpec and TypeVarTuple (#15626)
These are treated the same way as `TypeVar` except imported from `typing_extensions` instead of `typing` by default.
1 parent f72e4e5 commit 67cc059

File tree

2 files changed

+28
-6
lines changed

2 files changed

+28
-6
lines changed

mypy/stubgen.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -660,7 +660,7 @@ def visit_mypy_file(self, o: MypyFile) -> None:
660660
"_typeshed": ["Incomplete"],
661661
"typing": ["Any", "TypeVar", "NamedTuple"],
662662
"collections.abc": ["Generator"],
663-
"typing_extensions": ["TypedDict"],
663+
"typing_extensions": ["TypedDict", "ParamSpec", "TypeVarTuple"],
664664
}
665665
for pkg, imports in known_imports.items():
666666
for t in imports:
@@ -1158,10 +1158,14 @@ def is_alias_expression(self, expr: Expression, top_level: bool = True) -> bool:
11581158
Used to know if assignments look like type aliases, function alias,
11591159
or module alias.
11601160
"""
1161-
# Assignment of TypeVar(...) are passed through
1161+
# Assignment of TypeVar(...) and other typevar-likes are passed through
11621162
if isinstance(expr, CallExpr) and self.get_fullname(expr.callee) in (
11631163
"typing.TypeVar",
11641164
"typing_extensions.TypeVar",
1165+
"typing.ParamSpec",
1166+
"typing_extensions.ParamSpec",
1167+
"typing.TypeVarTuple",
1168+
"typing_extensions.TypeVarTuple",
11651169
):
11661170
return True
11671171
elif isinstance(expr, EllipsisExpr):

test-data/unit/stubgen.test

+22-4
Original file line numberDiff line numberDiff line change
@@ -1036,11 +1036,16 @@ y: C
10361036

10371037
[case testTypeVarPreserved]
10381038
tv = TypeVar('tv')
1039+
ps = ParamSpec('ps')
1040+
tvt = TypeVarTuple('tvt')
10391041

10401042
[out]
10411043
from typing import TypeVar
1044+
from typing_extensions import ParamSpec, TypeVarTuple
10421045

10431046
tv = TypeVar('tv')
1047+
ps = ParamSpec('ps')
1048+
tvt = TypeVarTuple('tvt')
10441049

10451050
[case testTypeVarArgsPreserved]
10461051
tv = TypeVar('tv', int, str)
@@ -1052,29 +1057,37 @@ tv = TypeVar('tv', int, str)
10521057

10531058
[case testTypeVarNamedArgsPreserved]
10541059
tv = TypeVar('tv', bound=bool, covariant=True)
1060+
ps = ParamSpec('ps', bound=bool, covariant=True)
10551061

10561062
[out]
10571063
from typing import TypeVar
1064+
from typing_extensions import ParamSpec
10581065

10591066
tv = TypeVar('tv', bound=bool, covariant=True)
1067+
ps = ParamSpec('ps', bound=bool, covariant=True)
10601068

10611069
[case TypeVarImportAlias]
1062-
from typing import TypeVar as t_TV
1063-
from typing_extensions import TypeVar as te_TV
1070+
from typing import TypeVar as t_TV, ParamSpec as t_PS
1071+
from typing_extensions import TypeVar as te_TV, TypeVarTuple as te_TVT
10641072
from x import TypeVar as x_TV
10651073

10661074
T = t_TV('T')
10671075
U = te_TV('U')
10681076
V = x_TV('V')
10691077

1078+
PS = t_PS('PS')
1079+
TVT = te_TVT('TVT')
1080+
10701081
[out]
10711082
from _typeshed import Incomplete
1072-
from typing import TypeVar as t_TV
1073-
from typing_extensions import TypeVar as te_TV
1083+
from typing import ParamSpec as t_PS, TypeVar as t_TV
1084+
from typing_extensions import TypeVar as te_TV, TypeVarTuple as te_TVT
10741085

10751086
T = t_TV('T')
10761087
U = te_TV('U')
10771088
V: Incomplete
1089+
PS = t_PS('PS')
1090+
TVT = te_TVT('TVT')
10781091

10791092
[case testTypeVarFromImportAlias]
10801093
import typing as t
@@ -1085,6 +1098,9 @@ T = t.TypeVar('T')
10851098
U = te.TypeVar('U')
10861099
V = x.TypeVar('V')
10871100

1101+
PS = t.ParamSpec('PS')
1102+
TVT = te.TypeVarTuple('TVT')
1103+
10881104
[out]
10891105
import typing as t
10901106
import typing_extensions as te
@@ -1093,6 +1109,8 @@ from _typeshed import Incomplete
10931109
T = t.TypeVar('T')
10941110
U = te.TypeVar('U')
10951111
V: Incomplete
1112+
PS = t.ParamSpec('PS')
1113+
TVT = te.TypeVarTuple('TVT')
10961114

10971115
[case testTypeAliasPreserved]
10981116
alias = str

0 commit comments

Comments
 (0)