[ty] Track different uses of legacy typevars, including context when rendering typevars (#19604)

This PR introduces a few related changes:

- We now keep track of each time a legacy typevar is bound in a
different generic context (e.g. class, function), and internally create
a new `TypeVarInstance` for each usage. This means the rest of the code
can now assume that salsa-equivalent `TypeVarInstance`s refer to the
same typevar, even taking into account that legacy typevars can be used
more than once.

- We also go ahead and track the binding context of PEP 695 typevars.
That's _much_ easier to track since we have the binding context right
there during type inference.

- With that in place, we can now include the name of the binding context
when rendering typevars (e.g. `T@f` instead of `T`)

crates/ty_python_semantic/src/semantic_index/definition.rs

@@ -60,6 +60,35 @@ impl<'db> Definition<'db> {
         FileRange::new(self.file(db), self.kind(db).target_range(module))
     }
 
+    /// Returns the name of the item being defined, if applicable.
+    pub fn name(self, db: &'db dyn Db) -> Option<String> {
+        let file = self.file(db);
+        let module = parsed_module(db, file).load(db);
+        let kind = self.kind(db);
+        match kind {
+            DefinitionKind::Function(def) => {
+                let node = def.node(&module);
+                Some(node.name.as_str().to_string())
+            }
+            DefinitionKind::Class(def) => {
+                let node = def.node(&module);
+                Some(node.name.as_str().to_string())
+            }
+            DefinitionKind::TypeAlias(def) => {
+                let node = def.node(&module);
+                Some(
+                    node.name
+                        .as_name_expr()
+                        .expect("type alias name should be a NameExpr")
+                        .id
+                        .as_str()
+                        .to_string(),
+                )
+            }
+            _ => None,
+        }
+    }
+
     /// Extract a docstring from this definition, if applicable.
     /// This method returns a docstring for function and class definitions.
     /// The docstring is extracted from the first statement in the body if it's a string literal.

crates/ty_python_semantic/src/types.rs

@@ -701,6 +701,7 @@ impl<'db> Type<'db> {
                     Name::new_static("T_all"),
                     None,
                     None,
+                    None,
                     variance,
                     None,
                     TypeVarKind::Pep695,
@@ -5128,6 +5129,7 @@ impl<'db> Type<'db> {
                     db,
                     Name::new(format!("{}'instance", typevar.name(db))),
                     None,
+                    None,
                     Some(bound_or_constraints),
                     typevar.variance(db),
                     None,
@@ -5295,10 +5297,12 @@ impl<'db> Type<'db> {
                     let instance = Type::ClassLiteral(class).to_instance(db).expect(
                         "nearest_enclosing_class must return type that can be instantiated",
                     );
+                    let class_definition = class.definition(db);
                     Ok(Type::TypeVar(TypeVarInstance::new(
                         db,
                         ast::name::Name::new_static("Self"),
-                        Some(class.definition(db)),
+                        Some(class_definition),
+                        Some(class_definition),
                         Some(TypeVarBoundOrConstraints::UpperBound(instance)),
                         TypeVarVariance::Invariant,
                         None,
@@ -5573,6 +5577,9 @@ impl<'db> Type<'db> {
                     partial.get(db, typevar).unwrap_or(self)
                 }
                 TypeMapping::PromoteLiterals => self,
+                TypeMapping::BindLegacyTypevars(binding_context) => {
+                    Type::TypeVar(typevar.with_binding_context(db, *binding_context))
+                }
             }
 
             Type::FunctionLiteral(function) => {
@@ -5660,7 +5667,8 @@ impl<'db> Type<'db> {
             | Type::BytesLiteral(_)
             | Type::EnumLiteral(_) => match type_mapping {
                 TypeMapping::Specialization(_) |
-                TypeMapping::PartialSpecialization(_) => self,
+                TypeMapping::PartialSpecialization(_) |
+                TypeMapping::BindLegacyTypevars(_) => self,
                 TypeMapping::PromoteLiterals => self.literal_fallback_instance(db)
                     .expect("literal type should have fallback instance type"),
             }
@@ -6009,6 +6017,9 @@ pub enum TypeMapping<'a, 'db> {
     /// Promotes any literal types to their corresponding instance types (e.g. `Literal["string"]`
     /// to `str`)
     PromoteLiterals,
+    /// Binds a legacy typevar with the generic context (class, function, type alias) that it is
+    /// being used in.
+    BindLegacyTypevars(Definition<'db>),
 }
 
 fn walk_type_mapping<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
@@ -6023,7 +6034,7 @@ fn walk_type_mapping<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
         TypeMapping::PartialSpecialization(specialization) => {
             walk_partial_specialization(db, specialization, visitor);
         }
-        TypeMapping::PromoteLiterals => {}
+        TypeMapping::PromoteLiterals | TypeMapping::BindLegacyTypevars(_) => {}
     }
 }
 
@@ -6037,6 +6048,9 @@ impl<'db> TypeMapping<'_, 'db> {
                 TypeMapping::PartialSpecialization(partial.to_owned())
             }
             TypeMapping::PromoteLiterals => TypeMapping::PromoteLiterals,
+            TypeMapping::BindLegacyTypevars(binding_context) => {
+                TypeMapping::BindLegacyTypevars(*binding_context)
+            }
         }
     }
 
@@ -6049,6 +6063,9 @@ impl<'db> TypeMapping<'_, 'db> {
                 TypeMapping::PartialSpecialization(partial.normalized_impl(db, visitor))
             }
             TypeMapping::PromoteLiterals => TypeMapping::PromoteLiterals,
+            TypeMapping::BindLegacyTypevars(binding_context) => {
+                TypeMapping::BindLegacyTypevars(*binding_context)
+            }
         }
     }
 }
@@ -6582,6 +6599,35 @@ pub struct TypeVarInstance<'db> {
     /// The type var's definition (None if synthesized)
     pub definition: Option<Definition<'db>>,
 
+    /// The definition of the generic class, function, or type alias that binds this typevar. This
+    /// is `None` for a legacy typevar outside of a context that can bind it.
+    ///
+    /// For a legacy typevar, the binding context might be missing:
+    ///
+    /// ```py
+    /// T = TypeVar("T")                       # [1]
+    /// def generic_function(t: T) -> T: ...   # [2]
+    /// ```
+    ///
+    /// Here, we will create two `TypeVarInstance`s for the typevar `T`. Both will have `[1]` as
+    /// their [`definition`][Self::definition]. The first represents the variable when it is first
+    /// created, and not yet used, so it's `binding_context` will be `None`. The second represents
+    /// when the typevar is used in `generic_function`, and its `binding_context` will be `[2]`
+    /// (that is, the definition of `generic_function`).
+    ///
+    /// For a PEP 695 typevar, there will always be a binding context, since you can only define
+    /// one as part of creating the generic context that uses it:
+    ///
+    /// ```py
+    /// def generic_function[T](t: T) -> T: ...
+    /// ```
+    ///
+    /// Here, we will create a single `TypeVarInstance`. Its [`definition`][Self::definition] will
+    /// be the `T` in `[T]` (i.e., the definition of the typevar in the syntactic construct that
+    /// creates the generic context that uses it). Its `binding_context` will be the definition of
+    /// `generic_function`.
+    binding_context: Option<Definition<'db>>,
+
     /// The upper bound or constraint on the type of this TypeVar
     bound_or_constraints: Option<TypeVarBoundOrConstraints<'db>>,
 
@@ -6611,6 +6657,25 @@ fn walk_type_var_type<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
 }
 
 impl<'db> TypeVarInstance<'db> {
+    pub(crate) fn with_binding_context(
+        self,
+        db: &'db dyn Db,
+        binding_context: Definition<'db>,
+    ) -> Self {
+        Self::new(
+            db,
+            self.name(db),
+            self.definition(db),
+            Some(binding_context),
+            self.bound_or_constraints(db),
+            self.variance(db),
+            self.default_ty(db).map(|ty| {
+                ty.apply_type_mapping(db, &TypeMapping::BindLegacyTypevars(binding_context))
+            }),
+            self.kind(db),
+        )
+    }
+
     pub(crate) fn is_legacy(self, db: &'db dyn Db) -> bool {
         matches!(self.kind(db), TypeVarKind::Legacy)
     }
@@ -6640,6 +6705,7 @@ impl<'db> TypeVarInstance<'db> {
             db,
             self.name(db),
             self.definition(db),
+            self.binding_context(db),
             self.bound_or_constraints(db)
                 .map(|b| b.normalized_impl(db, visitor)),
             self.variance(db),
@@ -6653,6 +6719,7 @@ impl<'db> TypeVarInstance<'db> {
             db,
             self.name(db),
             self.definition(db),
+            self.binding_context(db),
             self.bound_or_constraints(db)
                 .map(|b| b.materialize(db, variance)),
             self.variance(db),

crates/ty_python_semantic/src/types/class.rs

@@ -1273,7 +1273,21 @@ impl<'db> ClassLiteral<'db> {
         class_stmt
             .bases()
             .iter()
-            .map(|base_node| definition_expression_type(db, class_definition, base_node))
+            .map(
+                |base_node| match definition_expression_type(db, class_definition, base_node) {
+                    Type::KnownInstance(KnownInstanceType::SubscriptedGeneric(generic_context)) => {
+                        Type::KnownInstance(KnownInstanceType::SubscriptedGeneric(
+                            generic_context.with_binding_context(db, class_definition),
+                        ))
+                    }
+                    Type::KnownInstance(KnownInstanceType::SubscriptedProtocol(
+                        generic_context,
+                    )) => Type::KnownInstance(KnownInstanceType::SubscriptedProtocol(
+                        generic_context.with_binding_context(db, class_definition),
+                    )),
+                    ty => ty,
+                },
+            )
             .collect()
     }
 
@@ -4098,11 +4112,15 @@ impl KnownClass {
                 };
 
                 let containing_assignment = index.expect_single_definition(target);
+                // A freshly created legacy TypeVar does not have a binding context until it is
+                // used in a base class list, function parameter list, or type alias.
+                let binding_context = None;
                 overload.set_return_type(Type::KnownInstance(KnownInstanceType::TypeVar(
                     TypeVarInstance::new(
                         db,
                         &target.id,
                         Some(containing_assignment),
+                        binding_context,
                         bound_or_constraint,
                         variance,
                         *default,

crates/ty_python_semantic/src/types/display.rs

@@ -205,7 +205,16 @@ impl Display for DisplayRepresentation<'_> {
                 )
             }
             Type::Tuple(specialization) => specialization.tuple(self.db).display(self.db).fmt(f),
-            Type::TypeVar(typevar) => f.write_str(typevar.name(self.db)),
+            Type::TypeVar(typevar) => {
+                f.write_str(typevar.name(self.db))?;
+                if let Some(binding_context) = typevar
+                    .binding_context(self.db)
+                    .and_then(|def| def.name(self.db))
+                {
+                    write!(f, "@{binding_context}")?;
+                }
+                Ok(())
+            }
             Type::AlwaysTruthy => f.write_str("AlwaysTruthy"),
             Type::AlwaysFalsy => f.write_str("AlwaysFalsy"),
             Type::BoundSuper(bound_super) => {

crates/ty_python_semantic/src/types/generics.rs

@@ -14,13 +14,13 @@ use crate::types::signatures::{Parameter, Parameters, Signature};
 use crate::types::tuple::{TupleSpec, TupleType};
 use crate::types::{
     KnownInstanceType, Type, TypeMapping, TypeRelation, TypeTransformer, TypeVarBoundOrConstraints,
-    TypeVarInstance, TypeVarVariance, UnionType, binding_type, declaration_type,
+    TypeVarInstance, TypeVarKind, TypeVarVariance, UnionType, binding_type, declaration_type,
 };
 use crate::{Db, FxOrderSet};
 
 /// Returns an iterator of any generic context introduced by the given scope or any enclosing
 /// scope.
-fn enclosing_generic_contexts<'db>(
+pub(crate) fn enclosing_generic_contexts<'db>(
     db: &'db dyn Db,
     module: &ParsedModuleRef,
     index: &SemanticIndex<'db>,
@@ -175,6 +175,19 @@ impl<'db> GenericContext<'db> {
         Some(Self::new(db, variables))
     }
 
+    pub(crate) fn with_binding_context(
+        self,
+        db: &'db dyn Db,
+        binding_context: Definition<'db>,
+    ) -> Self {
+        let variables: FxOrderSet<_> = self
+            .variables(db)
+            .iter()
+            .map(|typevar| typevar.with_binding_context(db, binding_context))
+            .collect();
+        Self::new(db, variables)
+    }
+
     pub(crate) fn len(self, db: &'db dyn Db) -> usize {
         self.variables(db).len()
     }
@@ -241,6 +254,22 @@ impl<'db> GenericContext<'db> {
         self.variables(db).is_subset(other.variables(db))
     }
 
+    pub(crate) fn binds_legacy_typevar(
+        self,
+        db: &'db dyn Db,
+        typevar: TypeVarInstance<'db>,
+    ) -> Option<TypeVarInstance<'db>> {
+        assert!(typevar.kind(db) == TypeVarKind::Legacy);
+        let typevar_def = typevar.definition(db);
+        self.variables(db)
+            .iter()
+            .find(|self_typevar| {
+                self_typevar.kind(db) == TypeVarKind::Legacy
+                    && self_typevar.definition(db) == typevar_def
+            })
+            .copied()
+    }
+
     /// Creates a specialization of this generic context. Panics if the length of `types` does not
     /// match the number of typevars in the generic context. You must provide a specific type for
     /// each typevar; no defaults are used. (Use [`specialize_partial`](Self::specialize_partial)

crates/ty_python_semantic/src/types/infer.rs

@@ -110,7 +110,7 @@ use crate::types::enums::is_enum_class;
 use crate::types::function::{
     FunctionDecorators, FunctionLiteral, FunctionType, KnownFunction, OverloadLiteral,
 };
-use crate::types::generics::GenericContext;
+use crate::types::generics::{GenericContext, enclosing_generic_contexts};
 use crate::types::mro::MroErrorKind;
 use crate::types::signatures::{CallableSignature, Signature};
 use crate::types::tuple::{TupleSpec, TupleType};
@@ -803,6 +803,9 @@ pub(super) struct TypeInferenceBuilder<'db, 'ast> {
     /// [`check_overloaded_functions`]: TypeInferenceBuilder::check_overloaded_functions
     called_functions: FxHashSet<FunctionType<'db>>,
 
+    /// Whether we are in a context that binds unbound legacy typevars.
+    legacy_typevar_binding_context: Option<Definition<'db>>,
+
     /// The deferred state of inferring types of certain expressions within the region.
     ///
     /// This is different from [`InferenceRegion::Deferred`] which works on the entire definition
@@ -847,6 +850,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             expressions: FxHashMap::default(),
             bindings: VecMap::default(),
             declarations: VecMap::default(),
+            legacy_typevar_binding_context: None,
             deferred: VecSet::default(),
             cycle_fallback: false,
         }
@@ -1713,9 +1717,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         match definition.kind(self.db()) {
             DefinitionKind::Function(function) => {
-                self.infer_function_deferred(function.node(self.module()));
+                self.infer_function_deferred(definition, function.node(self.module()));
+            }
+            DefinitionKind::Class(class) => {
+                self.infer_class_deferred(definition, class.node(self.module()));
             }
-            DefinitionKind::Class(class) => self.infer_class_deferred(class.node(self.module())),
             _ => {}
         }
     }
@@ -2207,6 +2213,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         self.infer_type_parameters(type_params);
 
         if let Some(arguments) = class.arguments.as_deref() {
+            // Note: We do not install a new `legacy_typevar_binding_context`; since this class has
+            // PEP 695 typevars, it should not also bind any legacy typevars via inheriting from
+            // `typing.Generic` or `typing.Protocol`.
             let mut call_arguments =
                 CallArguments::from_arguments(self.db(), arguments, |argument, splatted_value| {
                     let ty = self.infer_expression(splatted_value);
@@ -2228,6 +2237,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             .as_deref()
             .expect("function type params scope without type params");
 
+        // Note: We do not install a new `legacy_typevar_binding_context`; since this function has
+        // PEP 695 typevars, it should not also bind any legacy typevars by referencing them in its
+        // parameter or return type annotations.
         self.infer_return_type_annotation(
             function.returns.as_deref(),
             DeferredExpressionState::None,
@@ -2590,11 +2602,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             if self.defer_annotations() {
                 self.deferred.insert(definition);
             } else {
+                let previous_legacy_typevar_binding_context =
+                    self.legacy_typevar_binding_context.replace(definition);
                 self.infer_return_type_annotation(
                     returns.as_deref(),
                     DeferredExpressionState::None,
                 );
                 self.infer_parameters(parameters);
+                self.legacy_typevar_binding_context = previous_legacy_typevar_binding_context;
             }
         }
 
@@ -3006,25 +3021,38 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             if self.in_stub() || class_node.bases().iter().any(contains_string_literal) {
                 self.deferred.insert(definition);
             } else {
+                let previous_legacy_typevar_binding_context =
+                    self.legacy_typevar_binding_context.replace(definition);
                 for base in class_node.bases() {
                     self.infer_expression(base);
                 }
+                self.legacy_typevar_binding_context = previous_legacy_typevar_binding_context;
             }
         }
     }
 
-    fn infer_function_deferred(&mut self, function: &ast::StmtFunctionDef) {
+    fn infer_function_deferred(
+        &mut self,
+        definition: Definition<'db>,
+        function: &ast::StmtFunctionDef,
+    ) {
+        let previous_legacy_typevar_binding_context =
+            self.legacy_typevar_binding_context.replace(definition);
         self.infer_return_type_annotation(
             function.returns.as_deref(),
             DeferredExpressionState::Deferred,
         );
         self.infer_parameters(function.parameters.as_ref());
+        self.legacy_typevar_binding_context = previous_legacy_typevar_binding_context;
     }
 
-    fn infer_class_deferred(&mut self, class: &ast::StmtClassDef) {
+    fn infer_class_deferred(&mut self, definition: Definition<'db>, class: &ast::StmtClassDef) {
+        let previous_legacy_typevar_binding_context =
+            self.legacy_typevar_binding_context.replace(definition);
         for base in class.bases() {
             self.infer_expression(base);
         }
+        self.legacy_typevar_binding_context = previous_legacy_typevar_binding_context;
     }
 
     fn infer_type_alias_definition(
@@ -3330,6 +3358,28 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             bound,
             default,
         } = node;
+
+        // Find the binding context for the PEP 695 typevars defined in this scope. The typevar
+        // scope should have a child containing the class, function, or type alias definition. Find
+        // that scope and use its definition as the binding context.
+        let typevar_scope = definition.file_scope(self.db());
+        let child_scopes = self.index.child_scopes(typevar_scope);
+        let binding_context = child_scopes
+            .filter_map(|(_, binding_scope)| match binding_scope.node() {
+                NodeWithScopeKind::Class(class) => {
+                    Some(DefinitionNodeKey::from(class.node(self.context.module())))
+                }
+                NodeWithScopeKind::Function(function) => Some(DefinitionNodeKey::from(
+                    function.node(self.context.module()),
+                )),
+                NodeWithScopeKind::TypeAlias(alias) => {
+                    Some(DefinitionNodeKey::from(alias.node(self.context.module())))
+                }
+                _ => None,
+            })
+            .map(|key| self.index.expect_single_definition(key))
+            .next();
+
         let bound_or_constraint = match bound.as_deref() {
             Some(expr @ ast::Expr::Tuple(ast::ExprTuple { elts, .. })) => {
                 if elts.len() < 2 {
@@ -3374,6 +3424,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             self.db(),
             &name.id,
             Some(definition),
+            binding_context,
             bound_or_constraint,
             TypeVarVariance::Invariant, // TODO: infer this
             default_ty,
@@ -6351,6 +6402,52 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             self.infer_place_load(PlaceExprRef::from(&expr), ast::ExprRef::Name(name_node));
 
         resolved
+            .map_type(|ty| {
+                // If the expression resolves to a legacy typevar, we will have the TypeVarInstance
+                // that was created when the typevar was created, which will not have an associated
+                // binding context. If this expression appears inside of a generic context that
+                // binds that typevar, we need to update the TypeVarInstance to include that
+                // binding context. To do that, we walk the enclosing scopes, looking for the
+                // nearest generic context that binds the typevar.
+                //
+                // If the legacy typevar is still unbound after that search, and we are in a
+                // context that binds unbound legacy typevars (i.e., the signature of a generic
+                // function), bind it with that context.
+                let find_legacy_typevar_binding = |typevar: TypeVarInstance<'db>| {
+                    enclosing_generic_contexts(
+                        self.db(),
+                        self.context.module(),
+                        self.index,
+                        self.scope().file_scope_id(self.db()),
+                    )
+                    .find_map(|enclosing_context| {
+                        enclosing_context.binds_legacy_typevar(self.db(), typevar)
+                    })
+                    .or_else(|| {
+                        self.legacy_typevar_binding_context
+                            .map(|legacy_typevar_binding_context| {
+                                typevar
+                                    .with_binding_context(self.db(), legacy_typevar_binding_context)
+                            })
+                    })
+                };
+
+                match ty {
+                    Type::TypeVar(typevar) if typevar.is_legacy(self.db()) => {
+                        find_legacy_typevar_binding(typevar)
+                            .map(Type::TypeVar)
+                            .unwrap_or(ty)
+                    }
+                    Type::KnownInstance(KnownInstanceType::TypeVar(typevar))
+                        if typevar.is_legacy(self.db()) =>
+                    {
+                        find_legacy_typevar_binding(typevar)
+                            .map(|typevar| Type::KnownInstance(KnownInstanceType::TypeVar(typevar)))
+                            .unwrap_or(ty)
+                    }
+                    _ => ty,
+                }
+            })
             // Not found in the module's explicitly declared global symbols?
             // Check the "implicit globals" such as `__doc__`, `__file__`, `__name__`, etc.
             // These are looked up as attributes on `types.ModuleType`.
@@ -8933,6 +9030,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             cycle_fallback,
 
             // builder only state
+            legacy_typevar_binding_context: _,
             deferred_state: _,
             called_functions: _,
             index: _,
@@ -8992,6 +9090,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             cycle_fallback,
 
             // builder only state
+            legacy_typevar_binding_context: _,
             deferred_state: _,
             called_functions: _,
             index: _,
@@ -9054,6 +9153,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             declarations: _,
 
             // Builder only state
+            legacy_typevar_binding_context: _,
             deferred_state: _,
             called_functions: _,
             index: _,

crates/ty_python_semantic/src/types/signatures.rs

@@ -1790,7 +1790,7 @@ mod tests {
             a_annotated_ty.unwrap().display(&db).to_string(),
             "Unknown | A | B"
         );
-        assert_eq!(b_annotated_ty.unwrap().display(&db).to_string(), "T");
+        assert_eq!(b_annotated_ty.unwrap().display(&db).to_string(), "T@f");
     }
 
     #[test]
@@ -1835,7 +1835,7 @@ mod tests {
         assert_eq!(b_name, "b");
         // Parameter resolution deferred:
         assert_eq!(a_annotated_ty.unwrap().display(&db).to_string(), "A | B");
-        assert_eq!(b_annotated_ty.unwrap().display(&db).to_string(), "T");
+        assert_eq!(b_annotated_ty.unwrap().display(&db).to_string(), "T@f");
     }
 
     #[test]

crates/ty_python_semantic/src/types/subclass_of.rs

@@ -93,6 +93,7 @@ impl<'db> SubclassOfType<'db> {
                         db,
                         Name::new_static("T_all"),
                         None,
+                        None,
                         Some(TypeVarBoundOrConstraints::UpperBound(
                             KnownClass::Type.to_instance(db),
                         )),