[ty] Restrict uncached raw signature access (#25866)

## Summary

This follows up on [the post-merge
review](#25761 (comment))
of #25761.

`FunctionLiteral::last_definition_raw_signature` bypasses the tracked
`FunctionType` query and therefore must only be called for a function
defined in the same module. After #25761, the method was visible
throughout the `types` module, making it easy for a future caller to
accidentally introduce a cross-module AST dependency and
over-invalidation.

This restores the method's privacy and routes the existing same-module
callers through a dedicated `same_module_uncached_raw_signature`
wrapper. The wrapper's name and documentation make the restriction
explicit, while cross-module callers continue to use the tracked query.

Author: charliermarsh

crates/ty_python_semantic/src/types/function.rs

@@ -871,7 +871,7 @@ impl<'db> FunctionLiteral<'db> {
     /// calling query is not in the same file as this function is defined in, then this will create
     /// a cross-module dependency directly on the full AST which will lead to cache
     /// over-invalidation.
-    pub(super) fn last_definition_raw_signature(
+    fn last_definition_raw_signature(
         self,
         db: &'db dyn Db,
         return_callable_typevar_scope: ReturnCallableTypeVarScope,
@@ -953,6 +953,23 @@ impl<'db> FunctionLiteral<'db> {
     }
 }
 
+/// ## Warning
+///
+/// This uses the semantic index to find the definition of the function. This means that if the
+/// calling query is not in the same file as this function is defined in, then this will create
+/// a cross-module dependency directly on the full AST which will lead to cache
+/// over-invalidation. Cross-module callers should use the tracked
+/// [`FunctionType::last_definition_raw_signature`] query instead.
+pub(super) fn same_module_uncached_raw_signature<'db>(
+    db: &'db dyn Db,
+    function: FunctionType<'db>,
+    return_callable_typevar_scope: ReturnCallableTypeVarScope,
+) -> Signature<'db> {
+    function
+        .literal(db)
+        .last_definition_raw_signature(db, return_callable_typevar_scope)
+}
+
 /// Indicates whether a method is explicitly or implicitly abstract.
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, salsa::Update, get_size2::GetSize)]
 pub(super) enum AbstractMethodKind {

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -69,6 +69,7 @@ use crate::types::diagnostic::{
 use crate::types::enums::{enum_ignored_names, is_enum_class_by_inheritance};
 use crate::types::function::{
     FunctionDecorators, FunctionType, KnownFunction, report_revealed_type,
+    same_module_uncached_raw_signature,
 };
 use crate::types::generics::{
     GenericContext, InferableTypeVars, Specialization, SpecializationBuilder, bind_typevar,
@@ -5109,13 +5110,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     // the expected type passed should be the "raw" type,
                     // i.e. type variables in the return type are non-inferable,
                     // and the return types of async functions are not wrapped in `CoroutineType[...]`.
-                    let return_ty = func
-                        .literal(self.db())
-                        .last_definition_raw_signature(
-                            self.db(),
-                            ReturnCallableTypeVarScope::Lexical,
-                        )
-                        .return_ty;
+                    let return_ty = same_module_uncached_raw_signature(
+                        self.db(),
+                        func,
+                        ReturnCallableTypeVarScope::Lexical,
+                    )
+                    .return_ty;
 
                     // For generator functions, the declared return type is e.g.
                     // `Generator[YieldType, SendType, ReturnType]`. The type context
@@ -8748,10 +8748,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             let _ = self.infer_optional_expression(value.as_deref(), TypeContext::default());
             return Type::unknown();
         };
-        let declared_return_ty = enclosing_function
-            .literal(self.db())
-            .last_definition_raw_signature(self.db(), ReturnCallableTypeVarScope::Public)
-            .return_ty;
+        let declared_return_ty = same_module_uncached_raw_signature(
+            self.db(),
+            enclosing_function,
+            ReturnCallableTypeVarScope::Public,
+        )
+        .return_ty;
         let return_type_span = enclosing_function.spans(self.db()).return_type;
 
         let Some(generator_type_params) = declared_return_ty.generator_types(self.db()) else {
@@ -8797,10 +8799,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             let _ = self.infer_expression(value, TypeContext::default());
             return Type::unknown();
         };
-        let annotated_return_ty = enclosing_function
-            .literal(self.db())
-            .last_definition_raw_signature(self.db(), ReturnCallableTypeVarScope::Public)
-            .return_ty;
+        let annotated_return_ty = same_module_uncached_raw_signature(
+            self.db(),
+            enclosing_function,
+            ReturnCallableTypeVarScope::Public,
+        )
+        .return_ty;
 
         let Some(outer_expected) = annotated_return_ty.generator_types(self.db()) else {
             let _ = self.infer_expression(value, TypeContext::default());

crates/ty_python_semantic/src/types/infer/builder/function.rs

@@ -14,6 +14,7 @@ use crate::{
         function::{
             FunctionBodyKind, FunctionDecorators, FunctionLiteral, FunctionType, KnownFunction,
             OverloadLiteral, function_body_kind, is_implicit_classmethod,
+            same_module_uncached_raw_signature,
         },
         generics::{enclosing_generic_contexts, typing_self},
         infer::{
@@ -80,18 +81,18 @@ impl<'db> ExpectedReturnType<'db> {
 
         let public = normalize(
             db,
-            function
-                .literal(db)
-                .last_definition_raw_signature(db, ReturnCallableTypeVarScope::Public)
+            same_module_uncached_raw_signature(db, function, ReturnCallableTypeVarScope::Public)
                 .return_ty,
         );
         let lexical = function_node.type_params.is_some().then(|| {
             normalize(
                 db,
-                function
-                    .literal(db)
-                    .last_definition_raw_signature(db, ReturnCallableTypeVarScope::Lexical)
-                    .return_ty,
+                same_module_uncached_raw_signature(
+                    db,
+                    function,
+                    ReturnCallableTypeVarScope::Lexical,
+                )
+                .return_ty,
             )
         });
 
@@ -169,10 +170,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
             let enclosing_function = nearest_enclosing_function(db, self.index, self.scope())
                 .expect("should be in a function body scope");
-            let declared_ty = enclosing_function
-                .literal(db)
-                .last_definition_raw_signature(db, ReturnCallableTypeVarScope::Public)
-                .return_ty;
+            let declared_ty = same_module_uncached_raw_signature(
+                db,
+                enclosing_function,
+                ReturnCallableTypeVarScope::Public,
+            )
+            .return_ty;
             let expected_return =
                 ExpectedReturnType::from_function(db, enclosing_function, function);
             let expected_ty = expected_return.public();