Skip to content

Commit

Permalink
Add a bit more support for PEP 612 in stubs.
Browse files Browse the repository at this point in the history
This is to unblock python/typeshed#6221.
Also see #786.

PiperOrigin-RevId: 407168815
  • Loading branch information
rchen152 committed Nov 2, 2021
1 parent c0ab8a9 commit 6e33205
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 16 deletions.
38 changes: 23 additions & 15 deletions pytype/pyi/definitions.py
Original file line number Diff line number Diff line change
Expand Up @@ -511,23 +511,31 @@ def _parameterized_type(self, base_type: Any, parameters):
raise ParseError("[..., ...] not supported")
return pytd.GenericType(base_type=base_type, parameters=(element_type,))
else:
parameters = tuple(pytd.AnythingType() if p is self.ELLIPSIS else p
for p in parameters)
processed_parameters = []
# We do not yet support PEP 612, Parameter Specification Variables.
# To avoid blocking typeshed from adopting this PEP, we convert new
# features to approximations that only use supported features.
for p in parameters:
if p is self.ELLIPSIS:
processed = pytd.AnythingType()
elif (p in self.param_specs and
self._matches_full_name(base_type, "typing.Generic")):
# Replacing a ParamSpec with a TypeVar isn't correct, but it'll work
# for simple cases in which the filled value is also a ParamSpec.
self.type_params.append(pytd.TypeParameter(p.name))
processed = p
elif (p in self.param_specs or
(isinstance(p, pytd.GenericType) and
self._matches_full_name(p, _CONCATENATE_TYPES))):
processed = pytd.AnythingType()
else:
processed = p
processed_parameters.append(processed)
parameters = tuple(processed_parameters)
if self._matches_named_type(base_type, _TUPLE_TYPES):
return pytdgen.heterogeneous_tuple(base_type, parameters)
elif self._matches_named_type(base_type, _CALLABLE_TYPES):
callable_parameters = []
for p in parameters:
# We do not yet support PEP 612, Parameter Specification Variables.
# To avoid blocking typeshed from adopting this PEP, we convert new
# features to Any.
if p in self.param_specs or (
isinstance(p, pytd.GenericType) and
self._matches_full_name(p, _CONCATENATE_TYPES)):
callable_parameters.append(pytd.AnythingType())
else:
callable_parameters.append(p)
return pytdgen.pytd_callable(base_type, tuple(callable_parameters))
return pytdgen.pytd_callable(base_type, parameters)
else:
assert parameters
return pytd.GenericType(base_type=base_type, parameters=parameters)
Expand Down Expand Up @@ -558,7 +566,7 @@ def resolve_type(self, name: Union[str, pytd_node.Node]) -> pytd.Type:
def new_type(
self,
name: Union[str, pytd_node.Node],
parameters: Optional[List[pytd_node.Node]] = None
parameters: Optional[List[pytd.Type]] = None
) -> pytd.Type:
"""Return the AST for a type.
Expand Down
33 changes: 32 additions & 1 deletion pytype/pyi/parser_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -2851,7 +2851,6 @@ def f(x: Callable[P, R]) -> Callable[P, Awaitable[R]]: ...
def f(x: Callable[..., R]) -> Callable[..., Awaitable[R]]: ...
""")

@test_base.skip("ParamSpec in custom generic classes not supported yet")
def test_custom_generic(self):
self.check("""
from typing import Callable, Generic, ParamSpec, TypeVar
Expand All @@ -2862,6 +2861,38 @@ def test_custom_generic(self):
class X(Generic[T, P]):
f: Callable[P, int]
x: T
""", """
from typing import Callable, Generic, TypeVar
P = TypeVar('P')
T = TypeVar('T')
class X(Generic[T, P]):
f: Callable[..., int]
x: T
""")

def test_use_custom_generic(self):
self.check("""
from typing import Callable, Generic, TypeVar
from typing_extensions import ParamSpec
_T = TypeVar('_T')
_P = ParamSpec('_P')
class Foo(Generic[_P, _T]): ...
def f(x: Callable[_P, _T]) -> Foo[_P, _T]: ...
""", """
from typing import Any, Callable, Generic, TypeVar
from typing_extensions import ParamSpec
_P = TypeVar('_P')
_T = TypeVar('_T')
class Foo(Generic[_P, _T]): ...
def f(x: Callable[..., _T]) -> Foo[Any, _T]: ...
""")

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

0 comments on commit 6e33205

Please sign in to comment.