[ty] Generalize nested types in inferred return types (astral-sh#29488)

## Summary

Generalize/promote types in inferred return types that refer to local
nested definitions.

For example, consider this unannotated decorator:
```python
 def register(event: str):
     def wrapper(func):
         # ...
         return func

     return wrapper
```

We do not want to infer `def wrapper(func) -> Unknown` as the return
type of `register`. This would not just expose internal details, it
would also incorrectly suggest that each call to `register` returns the
same function object. Instead, we generalize to an anonymous `(func) ->
Unknown` callable.

For nested classes, we cannot easily do the same, so we return `Unknown`
instead. In the future, to improve this, we might want to synthesize a
protocol that has the same attributes and methods as the nested class.

## Test strategy

Updated existing mdtests

Author: sharkdp

crates/ty_ide/src/inlay_hints.rs

@@ -1082,11 +1082,9 @@ mod tests {
             return 42
         "#,
         );
-        // TODO: Do not expose `def wrapper` in the return-type hint for `stringify`,
-        // and ideally show the return type of `str` for the returned callable.
         assert_snapshot!(test.inlay_hints(), @r#"
 
-        def stringify(function)[ -> def wrapper() -> Unknown]:
+        def stringify(function)[ -> () -> str]:
             def wrapper()[ -> str]:
                 return str([object=]function())
 
@@ -1098,15 +1096,15 @@ mod tests {
 
         ---------------------------------------------
         info[inlay-hint-location]: Inlay Hint Target
-          --> stdlib/ty_extensions/_internal.pyi:LL:1
+          --> stdlib/builtins.pyi:LL:7
            |
-        LL | Unknown: _SpecialForm
-           | ^^^^^^^
+        LL | class str(Sequence[str]):
+           |       ^^^
         info: Source
-          --> main2.py:LL:46
+          --> main2.py:LL:35
            |
-        LL | def stringify(function)[ -> def wrapper() -> Unknown]:
-           |                                              ^^^^^^^
+        LL | def stringify(function)[ -> () -> str]:
+           |                                   ^^^
 
         info[inlay-hint-location]: Inlay Hint Target
           --> stdlib/builtins.pyi:LL:7

crates/ty_python_semantic/resources/mdtest/function/return_type_inference.md

@@ -299,8 +299,7 @@ def returns_class():
 Inner1 = returns_class()
 Inner2 = returns_class()
 
-# TODO: This should be `bool` or `Literal[False]`, not `Literal[True]`
-reveal_type(Inner1 is Inner2)  # revealed: Literal[True]
+reveal_type(Inner1 is Inner2)  # revealed: bool
 
 # TODO: This should ideally be an error
 i: Inner1 = Inner2()
@@ -310,15 +309,13 @@ Therefore, the return type of `returns_class` cannot be `<class 'Inner'>`. Inste
 just treat it as `Unknown`.
 
 ```py
-# TODO: Should be `Unknown` or maybe `type`
-reveal_type(Inner1)  # revealed: <class 'Inner'>
+reveal_type(Inner1)  # revealed: Unknown
 ```
 
 Unfortunately, this means that we cannot infer precise types for attributes/methods on those types:
 
 ```py
-# TODO: Should be `Unknown`
-reveal_type(Inner1().attr)  # revealed: int
+reveal_type(Inner1().attr)  # revealed: Unknown
 ```
 
 If we wanted to improve this, we could return a synthetic protocol representing the structure of the
@@ -342,11 +339,9 @@ def returns_callable():
 callable1 = returns_callable()
 callable2 = returns_callable()
 
-# TODO: This should be `bool` or `Literal[False]`, not `Literal[True]`
-reveal_type(callable1 is callable2)  # revealed: Literal[True]
+reveal_type(callable1 is callable2)  # revealed: bool
 
-# TODO: should be `Callable[[], int]`
-reveal_type(callable1)  # revealed: def inner() -> int
+reveal_type(callable1)  # revealed: () -> int
 reveal_type(callable1())  # revealed: int
 ```
 
@@ -369,12 +364,11 @@ def returns_instance[T](x: T):
 def _(x: int):
     inner = returns_instance(x)
 
-    # TODO: `Inner` would probably be wrong here. Ideally, we would return a synthetic protocol
-    # with a type of `int` for the `attr` attribute.
-    reveal_type(inner)  # revealed: Inner
+    # TODO: Ideally, we would return a synthetic protocol with an `int` attribute.
+    reveal_type(inner)  # revealed: Unknown
 
     # TODO: should ideally be `int`
-    reveal_type(inner.attr)  # revealed: T@returns_instance
+    reveal_type(inner.attr)  # revealed: Unknown
 ```
 
 ## Generators

crates/ty_python_semantic/src/types.rs

@@ -9576,6 +9576,51 @@ impl<'db> Type<'db> {
         tcx: TypeContext<'db>,
         visitor: &ApplyTypeMappingVisitor<'_, 'db>,
     ) -> Type<'db> {
+        if let TypeMapping::GeneralizeLocalTypes(scope) = type_mapping {
+            let is_local = |definition: Definition<'db>| {
+                definition.program_file(db) == scope.program_file(db)
+                    && semantic_index(db, scope.program_file(db))
+                        .ancestor_scopes(definition.file_scope(db))
+                        .any(|(id, _)| id == scope.file_scope_id(db))
+            };
+
+            let class = match self {
+                Type::ClassLiteral(class) => Some(class),
+                Type::GenericAlias(alias) => Some(alias.origin(db).into()),
+                Type::NominalInstance(instance) => Some(instance.class_literal(db, visitor.env)),
+                Type::SubclassOf(subclass) => match subclass.subclass_of() {
+                    SubclassOfInner::Class(class) => Some(class.class_literal(db)),
+                    _ => None,
+                },
+                _ => None,
+            };
+            if class
+                .and_then(|class| class.definition(db))
+                .is_some_and(is_local)
+            {
+                // See doc comment for `TypeMapping::GeneralizeLocalTypes`: we might want to
+                // create a synthetic protocol that represents the local class, instead of
+                // returning `Unknown`.
+                return Type::unknown();
+            }
+
+            if let Type::FunctionLiteral(function) = self
+                && is_local(function.definition(db))
+            {
+                // Preserve the signature, but not the identity of a function object.
+                return visitor.visit(db, self, type_mapping, || {
+                    Type::Callable(CallableType::new(
+                        db,
+                        function
+                            .signature_with_inferred_return_type(db)
+                            .into_owned(),
+                        function.callable_type_kind(db),
+                    ))
+                    .apply_type_mapping_impl(db, type_mapping, tcx, visitor)
+                });
+            }
+        }
+
         // If we are binding `typing.Self`, and this type is what we are binding `Self` to, return
         // early. This is not just an optimization, it also prevents us from infinitely expanding
         // the type, if it's something that can contain a `Self` reference.
@@ -9916,6 +9961,7 @@ impl<'db> Type<'db> {
                 | TypeMapping::Materialize(_)
                 | TypeMapping::ReplaceParameterDefaults
                 | TypeMapping::EagerExpansion
+                | TypeMapping::GeneralizeLocalTypes(_)
                 | TypeMapping::RescopeReturnCallables(_)
                 | TypeMapping::Promote(PromotionMode::Off, _)
                 | TypeMapping::Promote(
@@ -9938,6 +9984,7 @@ impl<'db> Type<'db> {
                 | TypeMapping::Promote(..)
                 | TypeMapping::ReplaceParameterDefaults
                 | TypeMapping::EagerExpansion
+                | TypeMapping::GeneralizeLocalTypes(_)
                 | TypeMapping::RescopeReturnCallables(_) => self,
                 TypeMapping::Materialize(materialization_kind) => match materialization_kind {
                     MaterializationKind::Top => Type::object(),
@@ -11309,6 +11356,29 @@ pub enum TypeMapping<'a, 'db> {
 
     /// Updates any `Callable` types in a function signature return type to be generic if possible.
     RescopeReturnCallables(&'a FxHashMap<CallableType<'db>, CallableType<'db>>),
+
+    /// This type mapping is used in return type inference. It "promotes" certain types to supertypes
+    /// or to dynamic types in order to prevent nested definitions from leaking into the inferred
+    /// return type.
+    ///
+    /// For example, consider this unannotated decorator:
+    /// ```python
+    ///  def register(event: str):
+    ///      def wrapper(func):
+    ///          # ...
+    ///          return func
+    ///
+    ///      return wrapper
+    /// ```
+    ///
+    /// We do not want to infer `def wrapper(func) -> Unknown` as the return type of `register`. This would
+    /// not just expose internal details, it would also incorrectly suggest that each call to `register`
+    /// returns the same function object. Instead, we generalize to an anonymous `(func) -> Unknown` callable.
+    ///
+    /// For nested classes, we cannot easily do the same, so we return `Unknown` instead. In the future, to
+    /// improve this, we might want to synthesize a protocol that has the same attributes and methods as the
+    /// nested class.
+    GeneralizeLocalTypes(ScopeId<'db>),
 }
 
 impl<'db> TypeMapping<'_, 'db> {
@@ -11372,6 +11442,7 @@ impl<'db> TypeMapping<'_, 'db> {
             | TypeMapping::Materialize(_)
             | TypeMapping::ReplaceParameterDefaults
             | TypeMapping::EagerExpansion
+            | TypeMapping::GeneralizeLocalTypes(_)
             | TypeMapping::RescopeReturnCallables(_) => Some(context),
             TypeMapping::BindSelf(binding) => {
                 if binding.binding_context().is_some() {
@@ -11424,6 +11495,7 @@ impl<'db> TypeMapping<'_, 'db> {
             | TypeMapping::ReplaceSelf { .. }
             | TypeMapping::ReplaceParameterDefaults
             | TypeMapping::EagerExpansion
+            | TypeMapping::GeneralizeLocalTypes(_)
             | TypeMapping::RescopeReturnCallables(_) => self.clone(),
         }
     }

crates/ty_python_semantic/src/types/function.rs

@@ -1794,7 +1794,14 @@ impl<'db> FunctionType<'db> {
         let Some(class_definition) = index.class_definition_of_method(scope.file_scope_id(db))
         else {
             // For free functions, we simply use the type inferred from the function body.
-            return infer_complete_scope_types(db, scope).return_type();
+            return infer_complete_scope_types(db, scope)
+                .return_type()
+                .apply_type_mapping(
+                    db,
+                    &env,
+                    &TypeMapping::GeneralizeLocalTypes(scope),
+                    TypeContext::default(),
+                );
         };
         if self
             .literal(db)
@@ -1833,7 +1840,14 @@ impl<'db> FunctionType<'db> {
             return_type = UnionType::from_two_elements(db, &env, return_type, Type::unknown());
         }
 
-        return_type.apply_optional_specialization(db, self.class_specialization(db))
+        return_type
+            .apply_type_mapping(
+                db,
+                &env,
+                &TypeMapping::GeneralizeLocalTypes(scope),
+                TypeContext::default(),
+            )
+            .apply_optional_specialization(db, self.class_specialization(db))
     }
 
     #[salsa::tracked(

crates/ty_python_semantic/src/types/known_instance.rs

@@ -617,6 +617,7 @@ impl<'db> KnownInstanceType<'db> {
                 | TypeMapping::Materialize(_)
                 | TypeMapping::ReplaceParameterDefaults
                 | TypeMapping::EagerExpansion
+                | TypeMapping::GeneralizeLocalTypes(_)
                 | TypeMapping::RescopeReturnCallables(_)
                 | TypeMapping::ApplyRecursiveSubstitution(_) => Type::KnownInstance(self),
             },

crates/ty_python_semantic/src/types/typevar.rs

@@ -1482,6 +1482,7 @@ impl<'db> BoundTypeVarInstance<'db> {
             | TypeMapping::ReplaceParameterDefaults
             | TypeMapping::BindLegacyTypevars(_)
             | TypeMapping::EagerExpansion
+            | TypeMapping::GeneralizeLocalTypes(_)
             | TypeMapping::RescopeReturnCallables(_)
             | TypeMapping::ApplyRecursiveSubstitution(_) => Type::TypeVar(self),
             TypeMapping::Materialize(materialization_kind) => {