[red-knot] Add Type.definition method (#17153)

## Summary

This is a follow up to the goto type definition PR. Specifically, that
we want to avoid exposing too many semantic model internals publicly.

I want to get some feedback on the approach taken. I think it goes into
the right direction but I'm not super happy with it.
The basic idea is that we add a `Type::definition` method which does the
"goto type definition". The parts that I think make it awkward:

* We can't directly return `Definition` because we don't create a
`Definition` for modules (but we could?). Although I think it makes
sense to possibly have a more public wrapper type anyway?
* It doesn't handle unions and intersections. Mainly because not all
elements in an intersection may have a definition and we only want to
show a navigation target for intersections if there's only a single
positive element (besides maybe `Unknown`).


An alternative design or an addition to this design is to introduce a
`SemanticAnalysis(Db)` struct that has methods like
`type_definition(&self, type)` which explicitly exposes the methods we
want. I don't feel comfortable design this API yet because it's unclear
how fine granular it has to be (and if it is very fine granular,
directly using `Type` might be better after all)


## Test Plan

`cargo test`

crates/red_knot_ide/Cargo.toml

@@ -17,6 +17,7 @@ ruff_python_parser = { workspace = true }
 ruff_text_size = { workspace = true }
 red_knot_python_semantic = { workspace = true }
 
+rustc-hash = { workspace = true }
 salsa = { workspace = true }
 smallvec = { workspace = true }
 tracing = { workspace = true }

crates/red_knot_ide/src/goto.rs

@@ -24,9 +24,11 @@ pub fn goto_type_definition(
         ty.display(db.upcast())
     );
 
+    let navigation_targets = ty.navigation_targets(db);
+
     Some(RangedValue {
         range: FileRange::new(file, goto_target.range()),
-        value: ty.navigation_targets(db),
+        value: navigation_targets,
     })
 }
 
@@ -391,12 +393,12 @@ mod tests {
 
         test.write_file("lib.py", "a = 10").unwrap();
 
-        assert_snapshot!(test.goto_type_definition(), @r###"
+        assert_snapshot!(test.goto_type_definition(), @r"
         info: lint:goto-type-definition: Type definition
          --> /lib.py:1:1
           |
         1 | a = 10
-          | ^
+          | ^^^^^^
           |
         info: Source
          --> /main.py:4:13
@@ -406,7 +408,7 @@ mod tests {
         4 |             lib
           |             ^^^
           |
-        "###);
+        ");
     }
 
     #[test]
@@ -756,14 +758,13 @@ f(**kwargs<CURSOR>)
 
         assert_snapshot!(test.goto_type_definition(), @r"
         info: lint:goto-type-definition: Type definition
-           --> stdlib/builtins.pyi:443:7
+           --> stdlib/types.pyi:677:11
             |
-        441 |     def __getitem__(self, key: int, /) -> str | int | None: ...
-        442 |
-        443 | class str(Sequence[str]):
-            |       ^^^
-        444 |     @overload
-        445 |     def __new__(cls, object: object = ...) -> Self: ...
+        675 | if sys.version_info >= (3, 10):
+        676 |     @final
+        677 |     class NoneType:
+            |           ^^^^^^^^
+        678 |         def __bool__(self) -> Literal[False]: ...
             |
         info: Source
          --> /main.py:3:17
@@ -774,13 +775,14 @@ f(**kwargs<CURSOR>)
           |
 
         info: lint:goto-type-definition: Type definition
-           --> stdlib/types.pyi:677:11
+           --> stdlib/builtins.pyi:443:7
             |
-        675 | if sys.version_info >= (3, 10):
-        676 |     @final
-        677 |     class NoneType:
-            |           ^^^^^^^^
-        678 |         def __bool__(self) -> Literal[False]: ...
+        441 |     def __getitem__(self, key: int, /) -> str | int | None: ...
+        442 |
+        443 | class str(Sequence[str]):
+            |       ^^^
+        444 |     @overload
+        445 |     def __new__(cls, object: object = ...) -> Self: ...
             |
         info: Source
          --> /main.py:3:17

crates/red_knot_ide/src/lib.rs

@@ -4,19 +4,17 @@ mod goto;
 mod hover;
 mod markup;
 
-use std::ops::{Deref, DerefMut};
-
 pub use db::Db;
 pub use goto::goto_type_definition;
 pub use hover::hover;
 pub use markup::MarkupKind;
-use red_knot_python_semantic::types::{
-    Class, ClassBase, ClassLiteralType, FunctionType, InstanceType, IntersectionType,
-    KnownInstanceType, ModuleLiteralType, Type,
-};
+
+use rustc_hash::FxHashSet;
+use std::ops::{Deref, DerefMut};
+
+use red_knot_python_semantic::types::{Type, TypeDefinition};
 use ruff_db::files::{File, FileRange};
-use ruff_db::source::source_text;
-use ruff_text_size::{Ranged, TextLen, TextRange};
+use ruff_text_size::{Ranged, TextRange};
 
 /// Information associated with a text range.
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
@@ -58,7 +56,7 @@ where
 }
 
 /// Target to which the editor can navigate to.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct NavigationTarget {
     file: File,
 
@@ -99,6 +97,17 @@ impl NavigationTargets {
         Self(smallvec::SmallVec::new())
     }
 
+    fn unique(targets: impl IntoIterator<Item = NavigationTarget>) -> Self {
+        let unique: FxHashSet<_> = targets.into_iter().collect();
+        if unique.is_empty() {
+            Self::empty()
+        } else {
+            let mut targets = unique.into_iter().collect::<Vec<_>>();
+            targets.sort_by_key(|target| (target.file, target.focus_range.start()));
+            Self(targets.into())
+        }
+    }
+
     fn iter(&self) -> std::slice::Iter<'_, NavigationTarget> {
         self.0.iter()
     }
@@ -129,7 +138,7 @@ impl<'a> IntoIterator for &'a NavigationTargets {
 
 impl FromIterator<NavigationTarget> for NavigationTargets {
     fn from_iter<T: IntoIterator<Item = NavigationTarget>>(iter: T) -> Self {
-        Self(iter.into_iter().collect())
+        Self::unique(iter)
     }
 }
 
@@ -140,143 +149,50 @@ pub trait HasNavigationTargets {
 impl HasNavigationTargets for Type<'_> {
     fn navigation_targets(&self, db: &dyn Db) -> NavigationTargets {
         match self {
-            Type::BoundMethod(method) => method.function(db).navigation_targets(db),
-            Type::FunctionLiteral(function) => function.navigation_targets(db),
-            Type::ModuleLiteral(module) => module.navigation_targets(db),
             Type::Union(union) => union
                 .iter(db.upcast())
                 .flat_map(|target| target.navigation_targets(db))
                 .collect(),
-            Type::ClassLiteral(class) => class.navigation_targets(db),
-            Type::Instance(instance) => instance.navigation_targets(db),
-            Type::KnownInstance(instance) => instance.navigation_targets(db),
-            Type::SubclassOf(subclass_of_type) => match subclass_of_type.subclass_of() {
-                ClassBase::Class(class) => class.navigation_targets(db),
-                ClassBase::Dynamic(_) => NavigationTargets::empty(),
-            },
 
-            Type::StringLiteral(_)
-            | Type::BooleanLiteral(_)
-            | Type::LiteralString
-            | Type::IntLiteral(_)
-            | Type::BytesLiteral(_)
-            | Type::SliceLiteral(_)
-            | Type::MethodWrapper(_)
-            | Type::WrapperDescriptor(_)
-            | Type::PropertyInstance(_)
-            | Type::Tuple(_) => self.to_meta_type(db.upcast()).navigation_targets(db),
+            Type::Intersection(intersection) => {
+                // Only consider the positive elements because the negative elements are mainly from narrowing constraints.
+                let mut targets = intersection
+                    .iter_positive(db.upcast())
+                    .filter(|ty| !ty.is_unknown());
 
-            Type::TypeVar(var) => {
-                let definition = var.definition(db);
-                let full_range = definition.full_range(db.upcast());
+                let Some(first) = targets.next() else {
+                    return NavigationTargets::empty();
+                };
 
-                NavigationTargets::single(NavigationTarget {
-                    file: full_range.file(),
-                    focus_range: definition.focus_range(db.upcast()).range(),
-                    full_range: full_range.range(),
-                })
+                match targets.next() {
+                    Some(_) => {
+                        // If there are multiple types in the intersection, we can't navigate to a single one
+                        // because the type is the intersection of all those types.
+                        NavigationTargets::empty()
+                    }
+                    None => first.navigation_targets(db),
+                }
             }
 
-            Type::Intersection(intersection) => intersection.navigation_targets(db),
-
-            Type::Dynamic(_)
-            | Type::Never
-            | Type::Callable(_)
-            | Type::AlwaysTruthy
-            | Type::AlwaysFalsy => NavigationTargets::empty(),
+            ty => ty
+                .definition(db.upcast())
+                .map(|definition| definition.navigation_targets(db))
+                .unwrap_or_else(NavigationTargets::empty),
         }
     }
 }
 
-impl HasNavigationTargets for FunctionType<'_> {
+impl HasNavigationTargets for TypeDefinition<'_> {
     fn navigation_targets(&self, db: &dyn Db) -> NavigationTargets {
-        let function_range = self.focus_range(db.upcast());
+        let full_range = self.full_range(db.upcast());
         NavigationTargets::single(NavigationTarget {
-            file: function_range.file(),
-            focus_range: function_range.range(),
-            full_range: self.full_range(db.upcast()).range(),
+            file: full_range.file(),
+            focus_range: self.focus_range(db.upcast()).unwrap_or(full_range).range(),
+            full_range: full_range.range(),
         })
     }
 }
 
-impl HasNavigationTargets for Class<'_> {
-    fn navigation_targets(&self, db: &dyn Db) -> NavigationTargets {
-        let class_range = self.focus_range(db.upcast());
-        NavigationTargets::single(NavigationTarget {
-            file: class_range.file(),
-            focus_range: class_range.range(),
-            full_range: self.full_range(db.upcast()).range(),
-        })
-    }
-}
-
-impl HasNavigationTargets for ClassLiteralType<'_> {
-    fn navigation_targets(&self, db: &dyn Db) -> NavigationTargets {
-        self.class().navigation_targets(db)
-    }
-}
-
-impl HasNavigationTargets for InstanceType<'_> {
-    fn navigation_targets(&self, db: &dyn Db) -> NavigationTargets {
-        self.class().navigation_targets(db)
-    }
-}
-
-impl HasNavigationTargets for ModuleLiteralType<'_> {
-    fn navigation_targets(&self, db: &dyn Db) -> NavigationTargets {
-        let file = self.module(db).file();
-        let source = source_text(db.upcast(), file);
-
-        NavigationTargets::single(NavigationTarget {
-            file,
-            focus_range: TextRange::default(),
-            full_range: TextRange::up_to(source.text_len()),
-        })
-    }
-}
-
-impl HasNavigationTargets for KnownInstanceType<'_> {
-    fn navigation_targets(&self, db: &dyn Db) -> NavigationTargets {
-        match self {
-            KnownInstanceType::TypeVar(var) => {
-                let definition = var.definition(db);
-                let full_range = definition.full_range(db.upcast());
-
-                NavigationTargets::single(NavigationTarget {
-                    file: full_range.file(),
-                    focus_range: definition.focus_range(db.upcast()).range(),
-                    full_range: full_range.range(),
-                })
-            }
-
-            // TODO: Track the definition of `KnownInstance` and navigate to their definition.
-            _ => NavigationTargets::empty(),
-        }
-    }
-}
-
-impl HasNavigationTargets for IntersectionType<'_> {
-    fn navigation_targets(&self, db: &dyn Db) -> NavigationTargets {
-        // Only consider the positive elements because the negative elements are mainly from narrowing constraints.
-        let mut targets = self
-            .iter_positive(db.upcast())
-            .filter(|ty| !ty.is_unknown());
-
-        let Some(first) = targets.next() else {
-            return NavigationTargets::empty();
-        };
-
-        match targets.next() {
-            Some(_) => {
-                // If there are multiple types in the intersection, we can't navigate to a single one
-                // because the type is the intersection of all those types.
-                NavigationTargets::empty()
-            }
-            None => first.navigation_targets(db),
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use crate::db::tests::TestDb;

crates/red_knot_python_semantic/src/semantic_index/definition.rs

@@ -13,17 +13,13 @@ use crate::Db;
 
 /// A definition of a symbol.
 ///
-/// ## Module-local type
-/// This type should not be used as part of any cross-module API because
-/// it holds a reference to the AST node. Range-offset changes
-/// then propagate through all usages, and deserialization requires
-/// reparsing the entire module.
+/// ## ID stability
+/// The `Definition`'s ID is stable when the only field that change is its `kind` (AST node).
 ///
-/// E.g. don't use this type in:
-///
-/// * a return type of a cross-module query
-/// * a field of a type that is a return type of a cross-module query
-/// * an argument of a cross-module query
+/// The `Definition` changes when the `file`, `scope`, or `symbol` change. This can be
+/// because a new scope gets inserted before the `Definition` or a new symbol is inserted
+/// before this `Definition`. However, the ID can be considered stable and it is okay to use
+/// `Definition` in cross-module` salsa queries or as a field on other salsa tracked structs.
 #[salsa::tracked(debug)]
 pub struct Definition<'db> {
     /// The file in which the definition occurs.

crates/red_knot_python_semantic/src/types.rs

@@ -21,9 +21,9 @@ pub(crate) use self::infer::{
     infer_deferred_types, infer_definition_types, infer_expression_type, infer_expression_types,
     infer_scope_types,
 };
-pub use self::narrow::KnownConstraintFunction;
+pub(crate) use self::narrow::KnownConstraintFunction;
 pub(crate) use self::signatures::{CallableSignature, Signature, Signatures};
-pub use self::subclass_of::SubclassOfType;
+pub(crate) use self::subclass_of::SubclassOfType;
 use crate::module_name::ModuleName;
 use crate::module_resolver::{file_to_module, resolve_module, KnownModule};
 use crate::semantic_index::ast_ids::HasScopedExpressionId;
@@ -33,16 +33,14 @@ use crate::semantic_index::{imported_modules, semantic_index};
 use crate::suppression::check_suppressions;
 use crate::symbol::{imported_symbol, Boundness, Symbol, SymbolAndQualifiers};
 use crate::types::call::{Bindings, CallArgumentTypes};
-pub use crate::types::class_base::ClassBase;
+pub(crate) use crate::types::class_base::ClassBase;
 use crate::types::diagnostic::{INVALID_TYPE_FORM, UNSUPPORTED_BOOL_CONVERSION};
 use crate::types::infer::infer_unpack_types;
 use crate::types::mro::{Mro, MroError, MroIterator};
 pub(crate) use crate::types::narrow::infer_narrowing_constraint;
 use crate::types::signatures::{Parameter, ParameterForm, ParameterKind, Parameters};
 use crate::{Db, FxOrderSet, Module, Program};
-pub use class::Class;
-pub(crate) use class::KnownClass;
-pub use class::{ClassLiteralType, InstanceType, KnownInstanceType};
+pub(crate) use class::{Class, ClassLiteralType, InstanceType, KnownClass, KnownInstanceType};
 
 mod builder;
 mod call;
@@ -61,6 +59,7 @@ mod subclass_of;
 mod type_ordering;
 mod unpacker;
 
+mod definition;
 #[cfg(test)]
 mod property_tests;
 
@@ -227,6 +226,7 @@ macro_rules! todo_type {
     };
 }
 
+pub use crate::types::definition::TypeDefinition;
 pub(crate) use todo_type;
 
 /// Represents an instance of `builtins.property`.
@@ -3826,6 +3826,68 @@ impl<'db> Type<'db> {
             _ => KnownClass::Str.to_instance(db),
         }
     }
+
+    /// Returns where this type is defined.
+    ///
+    /// It's the foundation for the editor's "Go to type definition" feature
+    /// where the user clicks on a value and it takes them to where the value's type is defined.
+    ///
+    /// This method returns `None` for unions and intersections because how these
+    /// should be handled, especially when some variants don't have definitions, is
+    /// specific to the call site.
+    pub fn definition(&self, db: &'db dyn Db) -> Option<TypeDefinition<'db>> {
+        match self {
+            Self::BoundMethod(method) => {
+                Some(TypeDefinition::Function(method.function(db).definition(db)))
+            }
+            Self::FunctionLiteral(function) => {
+                Some(TypeDefinition::Function(function.definition(db)))
+            }
+            Self::ModuleLiteral(module) => Some(TypeDefinition::Module(module.module(db))),
+            Self::ClassLiteral(class_literal) => {
+                Some(TypeDefinition::Class(class_literal.class().definition(db)))
+            }
+            Self::Instance(instance) => {
+                Some(TypeDefinition::Class(instance.class().definition(db)))
+            }
+            Self::KnownInstance(instance) => match instance {
+                KnownInstanceType::TypeVar(var) => {
+                    Some(TypeDefinition::TypeVar(var.definition(db)))
+                }
+                KnownInstanceType::TypeAliasType(type_alias) => {
+                    Some(TypeDefinition::TypeAlias(type_alias.definition(db)))
+                }
+                _ => None,
+            },
+
+            Self::SubclassOf(subclass_of_type) => match subclass_of_type.subclass_of() {
+                ClassBase::Class(class) => Some(TypeDefinition::Class(class.definition(db))),
+                ClassBase::Dynamic(_) => None,
+            },
+
+            Self::StringLiteral(_)
+            | Self::BooleanLiteral(_)
+            | Self::LiteralString
+            | Self::IntLiteral(_)
+            | Self::BytesLiteral(_)
+            | Self::SliceLiteral(_)
+            | Self::MethodWrapper(_)
+            | Self::WrapperDescriptor(_)
+            | Self::PropertyInstance(_)
+            | Self::Tuple(_) => self.to_meta_type(db).definition(db),
+
+            Self::TypeVar(var) => Some(TypeDefinition::TypeVar(var.definition(db))),
+
+            Self::Union(_) | Self::Intersection(_) => None,
+
+            // These types have no definition
+            Self::Dynamic(_)
+            | Self::Never
+            | Self::Callable(_)
+            | Self::AlwaysTruthy
+            | Self::AlwaysFalsy => None,
+        }
+    }
 }
 
 impl<'db> From<&Type<'db>> for Type<'db> {
@@ -4717,7 +4779,7 @@ pub struct FunctionType<'db> {
 
 #[salsa::tracked]
 impl<'db> FunctionType<'db> {
-    pub fn has_known_decorator(self, db: &dyn Db, decorator: FunctionDecorators) -> bool {
+    pub(crate) fn has_known_decorator(self, db: &dyn Db, decorator: FunctionDecorators) -> bool {
         self.decorators(db).contains(decorator)
     }
 
@@ -4743,6 +4805,12 @@ impl<'db> FunctionType<'db> {
         )
     }
 
+    pub(crate) fn definition(self, db: &'db dyn Db) -> Definition<'db> {
+        let body_scope = self.body_scope(db);
+        let index = semantic_index(db, body_scope.file(db));
+        index.expect_single_definition(body_scope.node(db).expect_function())
+    }
+
     /// Typed externally-visible signature for this function.
     ///
     /// This is the signature as seen by external callers, possibly modified by decorators and/or
@@ -4756,7 +4824,7 @@ impl<'db> FunctionType<'db> {
     /// Were this not a salsa query, then the calling query
     /// would depend on the function's AST and rerun for every change in that file.
     #[salsa::tracked(return_ref)]
-    pub fn signature(self, db: &'db dyn Db) -> Signature<'db> {
+    pub(crate) fn signature(self, db: &'db dyn Db) -> Signature<'db> {
         let internal_signature = self.internal_signature(db);
 
         if self.has_known_decorator(db, FunctionDecorators::OVERLOAD) {
@@ -4779,12 +4847,11 @@ impl<'db> FunctionType<'db> {
     fn internal_signature(self, db: &'db dyn Db) -> Signature<'db> {
         let scope = self.body_scope(db);
         let function_stmt_node = scope.node(db).expect_function();
-        let definition =
-            semantic_index(db, scope.file(db)).expect_single_definition(function_stmt_node);
+        let definition = self.definition(db);
         Signature::from_function(db, definition, function_stmt_node)
     }
 
-    pub fn is_known(self, db: &'db dyn Db, known_function: KnownFunction) -> bool {
+    pub(crate) fn is_known(self, db: &'db dyn Db, known_function: KnownFunction) -> bool {
         self.known(db) == Some(known_function)
     }
 }
@@ -4904,7 +4971,7 @@ impl KnownFunction {
 pub struct BoundMethodType<'db> {
     /// The function that is being bound. Corresponds to the `__func__` attribute on a
     /// bound method object
-    pub function: FunctionType<'db>,
+    pub(crate) function: FunctionType<'db>,
     /// The instance on which this method has been called. Corresponds to the `__self__`
     /// attribute on a bound method object
     self_instance: Type<'db>,
@@ -5559,12 +5626,18 @@ pub struct TypeAliasType<'db> {
 
 #[salsa::tracked]
 impl<'db> TypeAliasType<'db> {
-    #[salsa::tracked]
-    pub fn value_type(self, db: &'db dyn Db) -> Type<'db> {
+    pub(crate) fn definition(self, db: &'db dyn Db) -> Definition<'db> {
         let scope = self.rhs_scope(db);
         let type_alias_stmt_node = scope.node(db).expect_type_alias();
-        let definition =
-            semantic_index(db, scope.file(db)).expect_single_definition(type_alias_stmt_node);
+
+        semantic_index(db, scope.file(db)).expect_single_definition(type_alias_stmt_node)
+    }
+
+    #[salsa::tracked]
+    pub(crate) fn value_type(self, db: &'db dyn Db) -> Type<'db> {
+        let scope = self.rhs_scope(db);
+        let type_alias_stmt_node = scope.node(db).expect_type_alias();
+        let definition = self.definition(db);
         definition_expression_type(db, definition, &type_alias_stmt_node.value)
     }
 }
@@ -5613,7 +5686,11 @@ impl<'db> UnionType<'db> {
         Self::from_elements(db, self.elements(db).iter().map(transform_fn))
     }
 
-    pub fn filter(&self, db: &'db dyn Db, filter_fn: impl FnMut(&&Type<'db>) -> bool) -> Type<'db> {
+    pub(crate) fn filter(
+        self,
+        db: &'db dyn Db,
+        filter_fn: impl FnMut(&&Type<'db>) -> bool,
+    ) -> Type<'db> {
         Self::from_elements(db, self.elements(db).iter().filter(filter_fn))
     }
 
@@ -5708,7 +5785,7 @@ impl<'db> UnionType<'db> {
         }
     }
 
-    pub fn is_fully_static(self, db: &'db dyn Db) -> bool {
+    pub(crate) fn is_fully_static(self, db: &'db dyn Db) -> bool {
         self.elements(db).iter().all(|ty| ty.is_fully_static(db))
     }
 
@@ -5716,7 +5793,7 @@ impl<'db> UnionType<'db> {
     ///
     /// See [`Type::normalized`] for more details.
     #[must_use]
-    pub fn normalized(self, db: &'db dyn Db) -> Self {
+    pub(crate) fn normalized(self, db: &'db dyn Db) -> Self {
         let mut new_elements: Vec<Type<'db>> = self
             .elements(db)
             .iter()
@@ -5727,7 +5804,7 @@ impl<'db> UnionType<'db> {
     }
 
     /// Return `true` if `self` represents the exact same set of possible runtime objects as `other`
-    pub fn is_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
+    pub(crate) fn is_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
         /// Inlined version of [`UnionType::is_fully_static`] to avoid having to lookup
         /// `self.elements` multiple times in the Salsa db in this single method.
         #[inline]
@@ -5765,7 +5842,7 @@ impl<'db> UnionType<'db> {
 
     /// Return `true` if `self` has exactly the same set of possible static materializations as `other`
     /// (if `self` represents the same set of possible sets of possible runtime objects as `other`)
-    pub fn is_gradual_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
+    pub(crate) fn is_gradual_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
         if self == other {
             return true;
         }
@@ -5818,7 +5895,7 @@ impl<'db> IntersectionType<'db> {
     ///
     /// See [`Type::normalized`] for more details.
     #[must_use]
-    pub fn normalized(self, db: &'db dyn Db) -> Self {
+    pub(crate) fn normalized(self, db: &'db dyn Db) -> Self {
         fn normalized_set<'db>(
             db: &'db dyn Db,
             elements: &FxOrderSet<Type<'db>>,
@@ -5837,13 +5914,13 @@ impl<'db> IntersectionType<'db> {
         )
     }
 
-    pub fn is_fully_static(self, db: &'db dyn Db) -> bool {
+    pub(crate) fn is_fully_static(self, db: &'db dyn Db) -> bool {
         self.positive(db).iter().all(|ty| ty.is_fully_static(db))
             && self.negative(db).iter().all(|ty| ty.is_fully_static(db))
     }
 
     /// Return `true` if `self` represents exactly the same set of possible runtime objects as `other`
-    pub fn is_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
+    pub(crate) fn is_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
         /// Inlined version of [`IntersectionType::is_fully_static`] to avoid having to lookup
         /// `positive` and `negative` multiple times in the Salsa db in this single method.
         #[inline]
@@ -5898,7 +5975,7 @@ impl<'db> IntersectionType<'db> {
 
     /// Return `true` if `self` has exactly the same set of possible static materializations as `other`
     /// (if `self` represents the same set of possible sets of possible runtime objects as `other`)
-    pub fn is_gradual_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
+    pub(crate) fn is_gradual_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
         if self == other {
             return true;
         }
@@ -6040,13 +6117,13 @@ pub struct StringLiteralType<'db> {
 
 impl<'db> StringLiteralType<'db> {
     /// The length of the string, as would be returned by Python's `len()`.
-    pub fn python_len(&self, db: &'db dyn Db) -> usize {
+    pub(crate) fn python_len(self, db: &'db dyn Db) -> usize {
         self.value(db).chars().count()
     }
 
     /// Return an iterator over each character in the string literal.
     /// as would be returned by Python's `iter()`.
-    pub fn iter_each_char(&self, db: &'db dyn Db) -> impl Iterator<Item = Self> {
+    pub(crate) fn iter_each_char(self, db: &'db dyn Db) -> impl Iterator<Item = Self> {
         self.value(db)
             .chars()
             .map(|c| StringLiteralType::new(db, c.to_string().as_str()))
@@ -6060,7 +6137,7 @@ pub struct BytesLiteralType<'db> {
 }
 
 impl<'db> BytesLiteralType<'db> {
-    pub fn python_len(&self, db: &'db dyn Db) -> usize {
+    pub(crate) fn python_len(self, db: &'db dyn Db) -> usize {
         self.value(db).len()
     }
 }
@@ -6084,7 +6161,7 @@ pub struct TupleType<'db> {
 }
 
 impl<'db> TupleType<'db> {
-    pub fn from_elements<T: Into<Type<'db>>>(
+    pub(crate) fn from_elements<T: Into<Type<'db>>>(
         db: &'db dyn Db,
         types: impl IntoIterator<Item = T>,
     ) -> Type<'db> {
@@ -6105,7 +6182,7 @@ impl<'db> TupleType<'db> {
     ///
     /// See [`Type::normalized`] for more details.
     #[must_use]
-    pub fn normalized(self, db: &'db dyn Db) -> Self {
+    pub(crate) fn normalized(self, db: &'db dyn Db) -> Self {
         let elements: Box<[Type<'db>]> = self
             .elements(db)
             .iter()
@@ -6114,7 +6191,7 @@ impl<'db> TupleType<'db> {
         TupleType::new(db, elements)
     }
 
-    pub fn is_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
+    pub(crate) fn is_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
         let self_elements = self.elements(db);
         let other_elements = other.elements(db);
         self_elements.len() == other_elements.len()
@@ -6124,7 +6201,7 @@ impl<'db> TupleType<'db> {
                 .all(|(self_ty, other_ty)| self_ty.is_equivalent_to(db, *other_ty))
     }
 
-    pub fn is_gradual_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
+    pub(crate) fn is_gradual_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
         let self_elements = self.elements(db);
         let other_elements = other.elements(db);
         self_elements.len() == other_elements.len()
@@ -6335,16 +6412,11 @@ pub(crate) mod tests {
                 | KnownFunction::IsGradualEquivalentTo => KnownModule::KnotExtensions,
             };
 
-            let function_body_scope = known_module_symbol(&db, module, function_name)
+            let function_definition = known_module_symbol(&db, module, function_name)
                 .symbol
                 .expect_type()
                 .expect_function_literal()
-                .body_scope(&db);
-
-            let function_node = function_body_scope.node(&db).expect_function();
-
-            let function_definition = semantic_index(&db, function_body_scope.file(&db))
-                .expect_single_definition(function_node);
+                .definition(&db);
 
             assert_eq!(
                 KnownFunction::try_from_definition_and_name(&db, function_definition, function_name),

crates/red_knot_python_semantic/src/types/class.rs

@@ -1,5 +1,11 @@
 use std::sync::{LazyLock, Mutex};
 
+use super::{
+    class_base::ClassBase, infer_expression_type, infer_unpack_types, IntersectionBuilder,
+    KnownFunction, Mro, MroError, MroIterator, SubclassOfType, Truthiness, Type, TypeAliasType,
+    TypeQualifiers, TypeVarInstance,
+};
+use crate::semantic_index::definition::Definition;
 use crate::{
     module_resolver::file_to_module,
     semantic_index::{
@@ -22,12 +28,6 @@ use ruff_db::files::{File, FileRange};
 use ruff_python_ast::{self as ast, PythonVersion};
 use rustc_hash::FxHashSet;
 
-use super::{
-    class_base::ClassBase, infer_expression_type, infer_unpack_types, IntersectionBuilder,
-    KnownFunction, Mro, MroError, MroIterator, SubclassOfType, Truthiness, Type, TypeAliasType,
-    TypeQualifiers, TypeVarInstance,
-};
-
 /// Representation of a runtime class object.
 ///
 /// Does not in itself represent a type,
@@ -43,52 +43,14 @@ pub struct Class<'db> {
     pub(crate) known: Option<KnownClass>,
 }
 
-fn explicit_bases_cycle_recover<'db>(
-    _db: &'db dyn Db,
-    _value: &[Type<'db>],
-    _count: u32,
-    _self: Class<'db>,
-) -> salsa::CycleRecoveryAction<Box<[Type<'db>]>> {
-    salsa::CycleRecoveryAction::Iterate
-}
-
-fn explicit_bases_cycle_initial<'db>(_db: &'db dyn Db, _self: Class<'db>) -> Box<[Type<'db>]> {
-    Box::default()
-}
-
-fn try_mro_cycle_recover<'db>(
-    _db: &'db dyn Db,
-    _value: &Result<Mro<'db>, MroError<'db>>,
-    _count: u32,
-    _self: Class<'db>,
-) -> salsa::CycleRecoveryAction<Result<Mro<'db>, MroError<'db>>> {
-    salsa::CycleRecoveryAction::Iterate
-}
-
-#[allow(clippy::unnecessary_wraps)]
-fn try_mro_cycle_initial<'db>(
-    db: &'db dyn Db,
-    self_: Class<'db>,
-) -> Result<Mro<'db>, MroError<'db>> {
-    Ok(Mro::from_error(db, self_))
-}
-
-#[allow(clippy::ref_option, clippy::trivially_copy_pass_by_ref)]
-fn inheritance_cycle_recover<'db>(
-    _db: &'db dyn Db,
-    _value: &Option<InheritanceCycle>,
-    _count: u32,
-    _self: Class<'db>,
-) -> salsa::CycleRecoveryAction<Option<InheritanceCycle>> {
-    salsa::CycleRecoveryAction::Iterate
-}
-
-fn inheritance_cycle_initial<'db>(_db: &'db dyn Db, _self: Class<'db>) -> Option<InheritanceCycle> {
-    None
-}
-
 #[salsa::tracked]
 impl<'db> Class<'db> {
+    pub(crate) fn definition(self, db: &'db dyn Db) -> Definition<'db> {
+        let scope = self.body_scope(db);
+        let index = semantic_index(db, scope.file(db));
+        index.expect_single_definition(scope.node(db).expect_class())
+    }
+
     /// Return `true` if this class represents `known_class`
     pub(crate) fn is_known(self, db: &'db dyn Db, known_class: KnownClass) -> bool {
         self.known(db) == Some(known_class)
@@ -239,8 +201,7 @@ impl<'db> Class<'db> {
             .find_keyword("metaclass")?
             .value;
 
-        let class_definition =
-            semantic_index(db, self.file(db)).expect_single_definition(class_stmt);
+        let class_definition = self.definition(db);
 
         Some(definition_expression_type(
             db,
@@ -740,6 +701,50 @@ impl<'db> Class<'db> {
     }
 }
 
+fn explicit_bases_cycle_recover<'db>(
+    _db: &'db dyn Db,
+    _value: &[Type<'db>],
+    _count: u32,
+    _self: Class<'db>,
+) -> salsa::CycleRecoveryAction<Box<[Type<'db>]>> {
+    salsa::CycleRecoveryAction::Iterate
+}
+
+fn explicit_bases_cycle_initial<'db>(_db: &'db dyn Db, _self: Class<'db>) -> Box<[Type<'db>]> {
+    Box::default()
+}
+
+fn try_mro_cycle_recover<'db>(
+    _db: &'db dyn Db,
+    _value: &Result<Mro<'db>, MroError<'db>>,
+    _count: u32,
+    _self: Class<'db>,
+) -> salsa::CycleRecoveryAction<Result<Mro<'db>, MroError<'db>>> {
+    salsa::CycleRecoveryAction::Iterate
+}
+
+#[allow(clippy::unnecessary_wraps)]
+fn try_mro_cycle_initial<'db>(
+    db: &'db dyn Db,
+    self_: Class<'db>,
+) -> Result<Mro<'db>, MroError<'db>> {
+    Ok(Mro::from_error(db, self_))
+}
+
+#[allow(clippy::ref_option, clippy::trivially_copy_pass_by_ref)]
+fn inheritance_cycle_recover<'db>(
+    _db: &'db dyn Db,
+    _value: &Option<InheritanceCycle>,
+    _count: u32,
+    _self: Class<'db>,
+) -> salsa::CycleRecoveryAction<Option<InheritanceCycle>> {
+    salsa::CycleRecoveryAction::Iterate
+}
+
+fn inheritance_cycle_initial<'db>(_db: &'db dyn Db, _self: Class<'db>) -> Option<InheritanceCycle> {
+    None
+}
+
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 pub(super) enum InheritanceCycle {
     /// The class is cyclically defined and is a participant in the cycle.
@@ -763,7 +768,7 @@ pub struct ClassLiteralType<'db> {
 }
 
 impl<'db> ClassLiteralType<'db> {
-    pub fn class(self) -> Class<'db> {
+    pub(super) fn class(self) -> Class<'db> {
         self.class
     }
 
@@ -789,7 +794,7 @@ pub struct InstanceType<'db> {
 }
 
 impl<'db> InstanceType<'db> {
-    pub fn class(self) -> Class<'db> {
+    pub(super) fn class(self) -> Class<'db> {
         self.class
     }
 

crates/red_knot_python_semantic/src/types/class_base.rs

@@ -8,7 +8,7 @@ use itertools::Either;
 /// all types that would be invalid to have as a class base are
 /// transformed into [`ClassBase::unknown`]
 #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, salsa::Update)]
-pub enum ClassBase<'db> {
+pub(crate) enum ClassBase<'db> {
     Dynamic(DynamicType),
     Class(Class<'db>),
 }
@@ -18,7 +18,7 @@ impl<'db> ClassBase<'db> {
         Self::Dynamic(DynamicType::Any)
     }
 
-    pub const fn unknown() -> Self {
+    pub(crate) const fn unknown() -> Self {
         Self::Dynamic(DynamicType::Unknown)
     }
 

crates/red_knot_python_semantic/src/types/definition.rs

@@ -0,0 +1,39 @@
+use crate::semantic_index::definition::Definition;
+use crate::{Db, Module};
+use ruff_db::files::FileRange;
+use ruff_db::source::source_text;
+use ruff_text_size::{TextLen, TextRange};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub enum TypeDefinition<'db> {
+    Module(Module),
+    Class(Definition<'db>),
+    Function(Definition<'db>),
+    TypeVar(Definition<'db>),
+    TypeAlias(Definition<'db>),
+}
+
+impl TypeDefinition<'_> {
+    pub fn focus_range(&self, db: &dyn Db) -> Option<FileRange> {
+        match self {
+            Self::Module(_) => None,
+            Self::Class(definition)
+            | Self::Function(definition)
+            | Self::TypeVar(definition)
+            | Self::TypeAlias(definition) => Some(definition.focus_range(db)),
+        }
+    }
+
+    pub fn full_range(&self, db: &dyn Db) -> FileRange {
+        match self {
+            Self::Module(module) => {
+                let source = source_text(db.upcast(), module.file());
+                FileRange::new(module.file(), TextRange::up_to(source.text_len()))
+            }
+            Self::Class(definition)
+            | Self::Function(definition)
+            | Self::TypeVar(definition)
+            | Self::TypeAlias(definition) => definition.full_range(db),
+        }
+    }
+}

crates/red_knot_python_semantic/src/types/subclass_of.rs

@@ -52,17 +52,17 @@ impl<'db> SubclassOfType<'db> {
     }
 
     /// Return the inner [`ClassBase`] value wrapped by this `SubclassOfType`.
-    pub const fn subclass_of(self) -> ClassBase<'db> {
+    pub(crate) const fn subclass_of(self) -> ClassBase<'db> {
         self.subclass_of
     }
 
-    pub const fn is_dynamic(self) -> bool {
+    pub(crate) const fn is_dynamic(self) -> bool {
         // Unpack `self` so that we're forced to update this method if any more fields are added in the future.
         let Self { subclass_of } = self;
         subclass_of.is_dynamic()
     }
 
-    pub const fn is_fully_static(self) -> bool {
+    pub(crate) const fn is_fully_static(self) -> bool {
         !self.is_dynamic()
     }