Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Polymorphic inference: basic support for variadic types #15879

Merged
merged 6 commits into from
Aug 18, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 10 additions & 54 deletions mypy/applytype.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,13 @@
from typing import Callable, Sequence

import mypy.subtypes
from mypy.expandtype import expand_type, expand_unpack_with_variables
from mypy.nodes import ARG_STAR, Context
from mypy.expandtype import expand_type
from mypy.nodes import Context
from mypy.types import (
AnyType,
CallableType,
Instance,
ParamSpecType,
PartialType,
TupleType,
Type,
TypeVarId,
TypeVarLikeType,
Expand All @@ -21,7 +19,6 @@
UnpackType,
get_proper_type,
)
from mypy.typevartuples import find_unpack_in_list, replace_starargs


def get_target_type(
Expand Down Expand Up @@ -107,6 +104,8 @@ def apply_generic_arguments(
if target_type is not None:
id_to_type[tvar.id] = target_type

# TODO: validate arg_kinds/arg_names for ParamSpec and TypeVarTuple replacements,
# not just type variable bounds above.
param_spec = callable.param_spec()
if param_spec is not None:
nt = id_to_type.get(param_spec.id)
Expand All @@ -122,55 +121,9 @@ def apply_generic_arguments(
# Apply arguments to argument types.
var_arg = callable.var_arg()
if var_arg is not None and isinstance(var_arg.typ, UnpackType):
star_index = callable.arg_kinds.index(ARG_STAR)
callable = callable.copy_modified(
arg_types=(
[expand_type(at, id_to_type) for at in callable.arg_types[:star_index]]
+ [callable.arg_types[star_index]]
+ [expand_type(at, id_to_type) for at in callable.arg_types[star_index + 1 :]]
)
)

unpacked_type = get_proper_type(var_arg.typ.type)
if isinstance(unpacked_type, TupleType):
# Assuming for now that because we convert prefixes to positional arguments,
# the first argument is always an unpack.
expanded_tuple = expand_type(unpacked_type, id_to_type)
if isinstance(expanded_tuple, TupleType):
# TODO: handle the case where the tuple has an unpack. This will
# hit an assert below.
expanded_unpack = find_unpack_in_list(expanded_tuple.items)
if expanded_unpack is not None:
callable = callable.copy_modified(
arg_types=(
callable.arg_types[:star_index]
+ [expanded_tuple]
+ callable.arg_types[star_index + 1 :]
)
)
else:
callable = replace_starargs(callable, expanded_tuple.items)
else:
# TODO: handle the case for if we get a variable length tuple.
assert False, f"mypy bug: unimplemented case, {expanded_tuple}"
elif isinstance(unpacked_type, TypeVarTupleType):
expanded_tvt = expand_unpack_with_variables(var_arg.typ, id_to_type)
if isinstance(expanded_tvt, list):
for t in expanded_tvt:
assert not isinstance(t, UnpackType)
callable = replace_starargs(callable, expanded_tvt)
else:
assert isinstance(expanded_tvt, Instance)
assert expanded_tvt.type.fullname == "builtins.tuple"
callable = callable.copy_modified(
arg_types=(
callable.arg_types[:star_index]
+ [expanded_tvt.args[0]]
+ callable.arg_types[star_index + 1 :]
)
)
else:
assert False, "mypy bug: unhandled case applying unpack"
callable = expand_type(callable, id_to_type)
assert isinstance(callable, CallableType)
return callable.copy_modified(variables=[tv for tv in tvars if tv.id not in id_to_type])
else:
callable = callable.copy_modified(
arg_types=[expand_type(at, id_to_type) for at in callable.arg_types]
Expand All @@ -183,6 +136,9 @@ def apply_generic_arguments(
type_guard = None

# The callable may retain some type vars if only some were applied.
# TODO: move apply_poly() logic from checkexpr.py here when new inference
# becomes universally used (i.e. in all passes + in unification).
# With this new logic we can actually *add* some new free variables.
remaining_tvars = [tv for tv in tvars if tv.id not in id_to_type]

return callable.copy_modified(
Expand Down
24 changes: 17 additions & 7 deletions mypy/checkexpr.py
Original file line number Diff line number Diff line change
Expand Up @@ -2312,11 +2312,15 @@ def check_argument_types(
]
actual_kinds = [nodes.ARG_STAR] + [nodes.ARG_POS] * (len(actuals) - 1)

assert isinstance(orig_callee_arg_type, TupleType)
assert orig_callee_arg_type.items
callee_arg_types = orig_callee_arg_type.items
# TODO: can we really assert this? What if formal is just plain Unpack[Ts]?
assert isinstance(orig_callee_arg_type, UnpackType)
assert isinstance(orig_callee_arg_type.type, ProperType) and isinstance(
orig_callee_arg_type.type, TupleType
)
assert orig_callee_arg_type.type.items
callee_arg_types = orig_callee_arg_type.type.items
callee_arg_kinds = [nodes.ARG_STAR] + [nodes.ARG_POS] * (
len(orig_callee_arg_type.items) - 1
len(orig_callee_arg_type.type.items) - 1
)
expanded_tuple = True

Expand Down Expand Up @@ -5792,8 +5796,9 @@ def visit_param_spec(self, t: ParamSpecType) -> Type:
return super().visit_param_spec(t)

def visit_type_var_tuple(self, t: TypeVarTupleType) -> Type:
# TODO: Support polymorphic apply for TypeVarTuple.
raise PolyTranslationError()
if t in self.poly_tvars and t not in self.bound_tvars:
raise PolyTranslationError()
return super().visit_type_var_tuple(t)

def visit_type_alias_type(self, t: TypeAliasType) -> Type:
if not t.args:
Expand Down Expand Up @@ -5827,7 +5832,6 @@ def visit_instance(self, t: Instance) -> Type:
return t.copy_modified(args=new_args)
# There is the same problem with callback protocols as with aliases
# (callback protocols are essentially more flexible aliases to callables).
# Note: consider supporting bindings in instances, e.g. LRUCache[[x: T], T].
if t.args and t.type.is_protocol and t.type.protocol_members == ["__call__"]:
if t.type in self.seen_aliases:
raise PolyTranslationError()
Expand Down Expand Up @@ -5862,6 +5866,12 @@ def __init__(self) -> None:
def visit_type_var(self, t: TypeVarType) -> bool:
return True

def visit_param_spec(self, t: ParamSpecType) -> bool:
return True

def visit_type_var_tuple(self, t: TypeVarTupleType) -> bool:
return True


def has_erased_component(t: Type | None) -> bool:
return t is not None and t.accept(HasErasedComponentsQuery())
Expand Down
Loading