Use assignment definition as typevar binding context

Author: sharkdp

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 `<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.

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,
         }
     }

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.

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, '_> {
@@ -723,9 +723,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(),
         );