From eee6f25f2e061f3b0b16a9ef8826bcf6fe4bbd73 Mon Sep 17 00:00:00 2001 From: David Peter Date: Mon, 24 Nov 2025 11:27:40 +0100 Subject: [PATCH] Use assignment definition as typevar binding context --- .../resources/mdtest/implicit_type_aliases.md | 47 +++++++++++++------ .../src/semantic_index/definition.rs | 6 +++ .../src/types/infer/builder.rs | 8 ++++ .../types/infer/builder/type_expression.rs | 27 ++++++++++- 4 files changed, 71 insertions(+), 17 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/implicit_type_aliases.md b/crates/ty_python_semantic/resources/mdtest/implicit_type_aliases.md index 6a3a9ad37d..2c51eacdf5 100644 --- a/crates/ty_python_semantic/resources/mdtest/implicit_type_aliases.md +++ b/crates/ty_python_semantic/resources/mdtest/implicit_type_aliases.md @@ -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 ``, … instead? (and similar for some others below) -reveal_type(MyList) # revealed: -reveal_type(MyDict) # revealed: +reveal_type(MyList) # revealed: +reveal_type(MyDict) # revealed: reveal_type(MyType) # revealed: GenericAlias -reveal_type(IntAndType) # revealed: -reveal_type(Pair) # revealed: -reveal_type(Sum) # revealed: +reveal_type(IntAndType) # revealed: +reveal_type(Pair) # revealed: +reveal_type(Sum) # revealed: reveal_type(ListOrTuple) # revealed: types.UnionType reveal_type(ListOrTupleLegacy) # revealed: types.UnionType reveal_type(MyCallable) # revealed: GenericAlias reveal_type(AnnotatedType) # revealed: +# 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: +reveal_type(DictStrTo) # revealed: 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: +reveal_type(MyOtherList) # revealed: 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. diff --git a/crates/ty_python_semantic/src/semantic_index/definition.rs b/crates/ty_python_semantic/src/semantic_index/definition.rs index 03acaa5335..2659e75493 100644 --- a/crates/ty_python_semantic/src/semantic_index/definition.rs +++ b/crates/ty_python_semantic/src/semantic_index/definition.rs @@ -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, } } diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index ce0a6e1d66..2203c9e537 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -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. diff --git a/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs b/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs index c65a747368..73b6124c25 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs @@ -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(), );