Skip to content

Commit

Permalink
stubgen: Support ParamSpec and TypeVarTuple (#15626)
Browse files Browse the repository at this point in the history
These are treated the same way as `TypeVar` except imported from
`typing_extensions` instead of `typing` by default.
  • Loading branch information
hamdanal committed Jul 9, 2023
1 parent f72e4e5 commit 67cc059
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 6 deletions.
8 changes: 6 additions & 2 deletions mypy/stubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -660,7 +660,7 @@ def visit_mypy_file(self, o: MypyFile) -> None:
"_typeshed": ["Incomplete"],
"typing": ["Any", "TypeVar", "NamedTuple"],
"collections.abc": ["Generator"],
"typing_extensions": ["TypedDict"],
"typing_extensions": ["TypedDict", "ParamSpec", "TypeVarTuple"],
}
for pkg, imports in known_imports.items():
for t in imports:
Expand Down Expand Up @@ -1158,10 +1158,14 @@ def is_alias_expression(self, expr: Expression, top_level: bool = True) -> bool:
Used to know if assignments look like type aliases, function alias,
or module alias.
"""
# Assignment of TypeVar(...) are passed through
# Assignment of TypeVar(...) and other typevar-likes are passed through
if isinstance(expr, CallExpr) and self.get_fullname(expr.callee) in (
"typing.TypeVar",
"typing_extensions.TypeVar",
"typing.ParamSpec",
"typing_extensions.ParamSpec",
"typing.TypeVarTuple",
"typing_extensions.TypeVarTuple",
):
return True
elif isinstance(expr, EllipsisExpr):
Expand Down
26 changes: 22 additions & 4 deletions test-data/unit/stubgen.test
Original file line number Diff line number Diff line change
Expand Up @@ -1036,11 +1036,16 @@ y: C

[case testTypeVarPreserved]
tv = TypeVar('tv')
ps = ParamSpec('ps')
tvt = TypeVarTuple('tvt')

[out]
from typing import TypeVar
from typing_extensions import ParamSpec, TypeVarTuple

tv = TypeVar('tv')
ps = ParamSpec('ps')
tvt = TypeVarTuple('tvt')

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

[case testTypeVarNamedArgsPreserved]
tv = TypeVar('tv', bound=bool, covariant=True)
ps = ParamSpec('ps', bound=bool, covariant=True)

[out]
from typing import TypeVar
from typing_extensions import ParamSpec

tv = TypeVar('tv', bound=bool, covariant=True)
ps = ParamSpec('ps', bound=bool, covariant=True)

[case TypeVarImportAlias]
from typing import TypeVar as t_TV
from typing_extensions import TypeVar as te_TV
from typing import TypeVar as t_TV, ParamSpec as t_PS
from typing_extensions import TypeVar as te_TV, TypeVarTuple as te_TVT
from x import TypeVar as x_TV

T = t_TV('T')
U = te_TV('U')
V = x_TV('V')

PS = t_PS('PS')
TVT = te_TVT('TVT')

[out]
from _typeshed import Incomplete
from typing import TypeVar as t_TV
from typing_extensions import TypeVar as te_TV
from typing import ParamSpec as t_PS, TypeVar as t_TV
from typing_extensions import TypeVar as te_TV, TypeVarTuple as te_TVT

T = t_TV('T')
U = te_TV('U')
V: Incomplete
PS = t_PS('PS')
TVT = te_TVT('TVT')

[case testTypeVarFromImportAlias]
import typing as t
Expand All @@ -1085,6 +1098,9 @@ T = t.TypeVar('T')
U = te.TypeVar('U')
V = x.TypeVar('V')

PS = t.ParamSpec('PS')
TVT = te.TypeVarTuple('TVT')

[out]
import typing as t
import typing_extensions as te
Expand All @@ -1093,6 +1109,8 @@ from _typeshed import Incomplete
T = t.TypeVar('T')
U = te.TypeVar('U')
V: Incomplete
PS = t.ParamSpec('PS')
TVT = te.TypeVarTuple('TVT')

[case testTypeAliasPreserved]
alias = str
Expand Down

0 comments on commit 67cc059

Please sign in to comment.