Skip to content

Commit

Permalink
Support self-types containing ParamSpec (#15903)
Browse files Browse the repository at this point in the history
Fixes #14968
Fixes #13911

The fix is simple, as I predicted on Discord, we simply should use
`get_all_type_vars()` instead of `get_type_vars()` (that specifically
returns only `TypeVarType`). I also use this opportunity to tidy-up code
in `bind_self()`, it should be now more readable, and much faster
(especially when compiled with mypyc).

cc @A5rocks

---------

Co-authored-by: Alex Waygood <[email protected]>
  • Loading branch information
ilevkivskyi and AlexWaygood authored Aug 19, 2023
1 parent 1db3eb3 commit d7d502e
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 20 deletions.
34 changes: 14 additions & 20 deletions mypy/typeops.py
Original file line number Diff line number Diff line change
Expand Up @@ -303,7 +303,7 @@ class B(A): pass
return cast(F, func)
self_param_type = get_proper_type(func.arg_types[0])

variables: Sequence[TypeVarLikeType] = []
variables: Sequence[TypeVarLikeType]
if func.variables and supported_self_type(self_param_type):
from mypy.infer import infer_type_arguments

Expand All @@ -312,46 +312,40 @@ class B(A): pass
original_type = erase_to_bound(self_param_type)
original_type = get_proper_type(original_type)

all_ids = func.type_var_ids()
# Find which of method type variables appear in the type of "self".
self_ids = {tv.id for tv in get_all_type_vars(self_param_type)}
self_vars = [tv for tv in func.variables if tv.id in self_ids]

# Solve for these type arguments using the actual class or instance type.
typeargs = infer_type_arguments(
func.variables, self_param_type, original_type, is_supertype=True
self_vars, self_param_type, original_type, is_supertype=True
)
if (
is_classmethod
# TODO: why do we need the extra guards here?
and any(isinstance(get_proper_type(t), UninhabitedType) for t in typeargs)
and isinstance(original_type, (Instance, TypeVarType, TupleType))
):
# In case we call a classmethod through an instance x, fallback to type(x)
# In case we call a classmethod through an instance x, fallback to type(x).
typeargs = infer_type_arguments(
func.variables, self_param_type, TypeType(original_type), is_supertype=True
self_vars, self_param_type, TypeType(original_type), is_supertype=True
)

ids = [tid for tid in all_ids if any(tid == t.id for t in get_type_vars(self_param_type))]

# Technically, some constrains might be unsolvable, make them <nothing>.
# Update the method signature with the solutions found.
# Technically, some constraints might be unsolvable, make them <nothing>.
to_apply = [t if t is not None else UninhabitedType() for t in typeargs]

def expand(target: Type) -> Type:
return expand_type(target, {id: to_apply[all_ids.index(id)] for id in ids})

arg_types = [expand(x) for x in func.arg_types[1:]]
ret_type = expand(func.ret_type)
variables = [v for v in func.variables if v.id not in ids]
func = expand_type(func, {tv.id: arg for tv, arg in zip(self_vars, to_apply)})
variables = [v for v in func.variables if v not in self_vars]
else:
arg_types = func.arg_types[1:]
ret_type = func.ret_type
variables = func.variables

original_type = get_proper_type(original_type)
if isinstance(original_type, CallableType) and original_type.is_type_obj():
original_type = TypeType.make_normalized(original_type.ret_type)
res = func.copy_modified(
arg_types=arg_types,
arg_types=func.arg_types[1:],
arg_kinds=func.arg_kinds[1:],
arg_names=func.arg_names[1:],
variables=variables,
ret_type=ret_type,
bound_args=[original_type],
)
return cast(F, res)
Expand Down
42 changes: 42 additions & 0 deletions test-data/unit/check-selftype.test
Original file line number Diff line number Diff line change
Expand Up @@ -1973,3 +1973,45 @@ class B(A):
reveal_type(self.x.extra) # N: Revealed type is "builtins.int"
reveal_type(self.xs[0].extra) # N: Revealed type is "builtins.int"
[builtins fixtures/list.pyi]

[case testSelfTypesWithParamSpecExtract]
from typing import Any, Callable, Generic, TypeVar
from typing_extensions import ParamSpec

P = ParamSpec("P")
F = TypeVar("F", bound=Callable[..., Any])
class Example(Generic[F]):
def __init__(self, fn: F) -> None:
...
def __call__(self: Example[Callable[P, Any]], *args: P.args, **kwargs: P.kwargs) -> None:
...

def test_fn(a: int, b: str) -> None:
...

example = Example(test_fn)
example() # E: Missing positional arguments "a", "b" in call to "__call__" of "Example"
example(1, "b") # OK
[builtins fixtures/list.pyi]

[case testSelfTypesWithParamSpecInfer]
from typing import TypeVar, Protocol, Type, Callable
from typing_extensions import ParamSpec

R = TypeVar("R", covariant=True)
P = ParamSpec("P")
class AsyncP(Protocol[P]):
def meth(self, *args: P.args, **kwargs: P.kwargs) -> None:
...

class Async:
@classmethod
def async_func(cls: Type[AsyncP[P]]) -> Callable[P, int]:
...

class Add(Async):
def meth(self, x: int, y: int) -> None: ...

reveal_type(Add.async_func()) # N: Revealed type is "def (x: builtins.int, y: builtins.int) -> builtins.int"
reveal_type(Add().async_func()) # N: Revealed type is "def (x: builtins.int, y: builtins.int) -> builtins.int"
[builtins fixtures/classmethod.pyi]

0 comments on commit d7d502e

Please sign in to comment.