Skip to content

Commit 6e33205

Browse files
committed
Add a bit more support for PEP 612 in stubs.
This is to unblock python/typeshed#6221. Also see #786. PiperOrigin-RevId: 407168815
1 parent c0ab8a9 commit 6e33205

File tree

2 files changed

+55
-16
lines changed

2 files changed

+55
-16
lines changed

pytype/pyi/definitions.py

Lines changed: 23 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -511,23 +511,31 @@ def _parameterized_type(self, base_type: Any, parameters):
511511
raise ParseError("[..., ...] not supported")
512512
return pytd.GenericType(base_type=base_type, parameters=(element_type,))
513513
else:
514-
parameters = tuple(pytd.AnythingType() if p is self.ELLIPSIS else p
515-
for p in parameters)
514+
processed_parameters = []
515+
# We do not yet support PEP 612, Parameter Specification Variables.
516+
# To avoid blocking typeshed from adopting this PEP, we convert new
517+
# features to approximations that only use supported features.
518+
for p in parameters:
519+
if p is self.ELLIPSIS:
520+
processed = pytd.AnythingType()
521+
elif (p in self.param_specs and
522+
self._matches_full_name(base_type, "typing.Generic")):
523+
# Replacing a ParamSpec with a TypeVar isn't correct, but it'll work
524+
# for simple cases in which the filled value is also a ParamSpec.
525+
self.type_params.append(pytd.TypeParameter(p.name))
526+
processed = p
527+
elif (p in self.param_specs or
528+
(isinstance(p, pytd.GenericType) and
529+
self._matches_full_name(p, _CONCATENATE_TYPES))):
530+
processed = pytd.AnythingType()
531+
else:
532+
processed = p
533+
processed_parameters.append(processed)
534+
parameters = tuple(processed_parameters)
516535
if self._matches_named_type(base_type, _TUPLE_TYPES):
517536
return pytdgen.heterogeneous_tuple(base_type, parameters)
518537
elif self._matches_named_type(base_type, _CALLABLE_TYPES):
519-
callable_parameters = []
520-
for p in parameters:
521-
# We do not yet support PEP 612, Parameter Specification Variables.
522-
# To avoid blocking typeshed from adopting this PEP, we convert new
523-
# features to Any.
524-
if p in self.param_specs or (
525-
isinstance(p, pytd.GenericType) and
526-
self._matches_full_name(p, _CONCATENATE_TYPES)):
527-
callable_parameters.append(pytd.AnythingType())
528-
else:
529-
callable_parameters.append(p)
530-
return pytdgen.pytd_callable(base_type, tuple(callable_parameters))
538+
return pytdgen.pytd_callable(base_type, parameters)
531539
else:
532540
assert parameters
533541
return pytd.GenericType(base_type=base_type, parameters=parameters)
@@ -558,7 +566,7 @@ def resolve_type(self, name: Union[str, pytd_node.Node]) -> pytd.Type:
558566
def new_type(
559567
self,
560568
name: Union[str, pytd_node.Node],
561-
parameters: Optional[List[pytd_node.Node]] = None
569+
parameters: Optional[List[pytd.Type]] = None
562570
) -> pytd.Type:
563571
"""Return the AST for a type.
564572

pytype/pyi/parser_test.py

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2851,7 +2851,6 @@ def f(x: Callable[P, R]) -> Callable[P, Awaitable[R]]: ...
28512851
def f(x: Callable[..., R]) -> Callable[..., Awaitable[R]]: ...
28522852
""")
28532853

2854-
@test_base.skip("ParamSpec in custom generic classes not supported yet")
28552854
def test_custom_generic(self):
28562855
self.check("""
28572856
from typing import Callable, Generic, ParamSpec, TypeVar
@@ -2862,6 +2861,38 @@ def test_custom_generic(self):
28622861
class X(Generic[T, P]):
28632862
f: Callable[P, int]
28642863
x: T
2864+
""", """
2865+
from typing import Callable, Generic, TypeVar
2866+
2867+
P = TypeVar('P')
2868+
T = TypeVar('T')
2869+
2870+
class X(Generic[T, P]):
2871+
f: Callable[..., int]
2872+
x: T
2873+
""")
2874+
2875+
def test_use_custom_generic(self):
2876+
self.check("""
2877+
from typing import Callable, Generic, TypeVar
2878+
from typing_extensions import ParamSpec
2879+
2880+
_T = TypeVar('_T')
2881+
_P = ParamSpec('_P')
2882+
2883+
class Foo(Generic[_P, _T]): ...
2884+
2885+
def f(x: Callable[_P, _T]) -> Foo[_P, _T]: ...
2886+
""", """
2887+
from typing import Any, Callable, Generic, TypeVar
2888+
from typing_extensions import ParamSpec
2889+
2890+
_P = TypeVar('_P')
2891+
_T = TypeVar('_T')
2892+
2893+
class Foo(Generic[_P, _T]): ...
2894+
2895+
def f(x: Callable[..., _T]) -> Foo[Any, _T]: ...
28652896
""")
28662897

28672898
@test_base.skip("ParamSpec in custom generic classes not supported yet")

0 commit comments

Comments
 (0)