Use assignment definition as typevar binding context

This commit is contained in:
David Peter
2025-11-24 11:27:40 +01:00
parent 013d43a2dd
commit eee6f25f2e
4 changed files with 71 additions and 17 deletions

View File

@@ -191,13 +191,13 @@ def _(
reveal_type(int_or_callable) # revealed: int | ((str, /) -> bytes)
reveal_type(callable_or_int) # revealed: ((str, /) -> bytes) | int
# TODO should be Unknown | int
reveal_type(type_var_or_int) # revealed: typing.TypeVar | int
reveal_type(type_var_or_int) # revealed: T@TypeVarOrInt | int
# TODO should be int | Unknown
reveal_type(int_or_type_var) # revealed: int | typing.TypeVar
reveal_type(int_or_type_var) # revealed: int | T@IntOrTypeVar
# TODO should be Unknown | None
reveal_type(type_var_or_none) # revealed: typing.TypeVar | None
reveal_type(type_var_or_none) # revealed: T@TypeVarOrNone | None
# TODO should be None | Unknown
reveal_type(none_or_type_var) # revealed: None | typing.TypeVar
reveal_type(none_or_type_var) # revealed: None | T@NoneOrTypeVar
```
If a type is unioned with itself in a value expression, the result is just that type. No
@@ -391,17 +391,17 @@ AnnotatedType = Annotated[T, "tag"]
TransparentAlias = T
MyOptional = T | None
# TODO: Consider displaying this as `<class 'list[T]'>`, … instead? (and similar for some others below)
reveal_type(MyList) # revealed: <class 'list[typing.TypeVar]'>
reveal_type(MyDict) # revealed: <class 'dict[typing.TypeVar, typing.TypeVar]'>
reveal_type(MyList) # revealed: <class 'list[T@MyList]'>
reveal_type(MyDict) # revealed: <class 'dict[T@MyDict, U@MyDict]'>
reveal_type(MyType) # revealed: GenericAlias
reveal_type(IntAndType) # revealed: <class 'tuple[int, typing.TypeVar]'>
reveal_type(Pair) # revealed: <class 'tuple[typing.TypeVar, typing.TypeVar]'>
reveal_type(Sum) # revealed: <class 'tuple[typing.TypeVar, typing.TypeVar]'>
reveal_type(IntAndType) # revealed: <class 'tuple[int, T@IntAndType]'>
reveal_type(Pair) # revealed: <class 'tuple[T@Pair, T@Pair]'>
reveal_type(Sum) # revealed: <class 'tuple[T@Sum, U@Sum]'>
reveal_type(ListOrTuple) # revealed: types.UnionType
reveal_type(ListOrTupleLegacy) # revealed: types.UnionType
reveal_type(MyCallable) # revealed: GenericAlias
reveal_type(AnnotatedType) # revealed: <typing.Annotated special form>
# TODO: This should ideally be `T@TransparentAlias`
reveal_type(TransparentAlias) # revealed: typing.TypeVar
reveal_type(MyOptional) # revealed: types.UnionType
@@ -445,7 +445,7 @@ U = TypeVar("U")
DictStrTo = MyDict[str, U]
reveal_type(DictStrTo) # revealed: <class 'dict[str, typing.TypeVar]'>
reveal_type(DictStrTo) # revealed: <class 'dict[str, U@DictStrTo]'>
def _(
dict_str_to_int: DictStrTo[int],
@@ -480,7 +480,7 @@ A generic implicit type alias can also be used in another generic implicit type
```py
MyOtherList = MyList[T]
reveal_type(MyOtherList) # revealed: <class 'list[typing.TypeVar]'>
reveal_type(MyOtherList) # revealed: <class 'list[T@MyOtherList]'>
def _(
list_of_ints: MyOtherList[int],
@@ -498,11 +498,11 @@ def _(
my_callable: MyCallable,
):
# TODO: Should be `list[Unknown]`
reveal_type(my_list) # revealed: list[typing.TypeVar]
reveal_type(my_list) # revealed: list[T@MyList]
# TODO: Should be `dict[Unknown, Unknown]`
reveal_type(my_dict) # revealed: dict[typing.TypeVar, typing.TypeVar]
reveal_type(my_dict) # revealed: dict[T@MyDict, U@MyDict]
# TODO: Should be `(...) -> Unknown`
reveal_type(my_callable) # revealed: (...) -> typing.TypeVar
reveal_type(my_callable) # revealed: (...) -> T@MyCallable
```
(Generic) implicit type aliases can be used as base classes:
@@ -552,6 +552,23 @@ def _(
reveal_type(dict_too_few_args) # revealed: Unknown
```
Trying to specialize a non-name node results in an error:
```py
from ty_extensions import TypeOf
IntOrStr = int | str
def this_does_not_work() -> TypeOf[IntOrStr]:
raise NotImplementedError()
def _(
# error: [invalid-type-form] "Cannot specialize a non-name node in a type expression"
specialized: this_does_not_work()[int],
):
reveal_type(specialized) # revealed: Unknown
```
## `Literal`s
We also support `typing.Literal` in implicit type aliases.

View File

@@ -90,6 +90,12 @@ impl<'db> Definition<'db> {
.to_string(),
)
}
DefinitionKind::Assignment(assignment) => {
let target_node = assignment.target.node(&module);
target_node
.as_name_expr()
.map(|name_expr| name_expr.id.as_str().to_string())
}
_ => None,
}
}

View File

@@ -4739,6 +4739,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
unpacked.expression_type(target)
}
TargetKind::Single => {
// This could be an implicit type alias (OptionalList = list[T] | None). Use the definition
// of `OptionalList` as the typevar binding context while inferring the RHS (`list[T] | None`),
// in order to bind `T@OptionalList`.
let previous_typevar_binding_context =
self.typevar_binding_context.replace(definition);
let value_ty = if let Some(standalone_expression) = self.index.try_expression(value)
{
self.infer_standalone_expression_impl(value, standalone_expression, tcx)
@@ -4777,6 +4783,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
self.infer_expression(value, tcx)
};
self.typevar_binding_context = previous_typevar_binding_context;
// `TYPE_CHECKING` is a special variable that should only be assigned `False`
// at runtime, but is always considered `True` in type checking.
// See mdtest/known_constants.md#user-defined-type_checking for details.

View File

@@ -2,7 +2,6 @@ use itertools::Either;
use ruff_python_ast as ast;
use super::{DeferredExpressionState, TypeInferenceBuilder};
use crate::FxOrderSet;
use crate::types::diagnostic::{
self, INVALID_TYPE_FORM, NON_SUBSCRIPTABLE, report_invalid_argument_number_to_special_form,
report_invalid_arguments_to_callable,
@@ -16,6 +15,7 @@ use crate::types::{
KnownInstanceType, LintDiagnosticGuard, Parameter, Parameters, SpecialFormType, SubclassOfType,
Type, TypeAliasType, TypeContext, TypeIsType, TypeMapping, UnionBuilder, UnionType, todo_type,
};
use crate::{FxOrderSet, ResolvedDefinition, definitions_for_name};
/// Type expressions
impl<'db> TypeInferenceBuilder<'db, '_> {
@@ -759,9 +759,32 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
) -> Type<'db> {
let db = self.db();
let Some(value) = subscript.value.as_name_expr() else {
if let Some(builder) = self.context.report_lint(&INVALID_TYPE_FORM, subscript) {
builder.into_diagnostic("Cannot specialize a non-name node in a type expression");
}
return Type::unknown();
};
// TODO: This is an expensive call to an API that was never meant to be called from
// type inference. We plan to rework how `in_type_expression` works in the future.
// This new approach will make this call unnecessary, so for now, we accept the hit
// in performance.
let definitions = definitions_for_name(self.db(), self.file(), value);
let Some(type_alias_definition) =
definitions.iter().find_map(ResolvedDefinition::definition)
else {
if let Some(builder) = self.context.report_lint(&INVALID_TYPE_FORM, subscript) {
builder.into_diagnostic(
"Cannot specialize implicit type alias with unknown definition",
);
}
return Type::unknown();
};
let generic_type_alias = value_ty.apply_type_mapping(
db,
&TypeMapping::BindLegacyTypevars(BindingContext::Synthetic),
&TypeMapping::BindLegacyTypevars(BindingContext::Definition(type_alias_definition)),
TypeContext::default(),
);