[red-knot] infer function's return type

Author: mtshiba

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

@@ -182,6 +182,43 @@ def f(cond: bool) -> int:
         return 2
 ```
 
+## Inferred return type
+
+If a function's return type is not annotated, it is inferred. The inferred type is the union of all
+possible return types.
+
+```py
+def f():
+    return 1
+
+reveal_type(f())  # revealed: Literal[1]
+
+def g(cond: bool):
+    if cond:
+        return 1
+    else:
+        return "a"
+
+reveal_type(g(True))  # revealed: Literal[1, "a"]
+
+# This function implicitly returns `None`.
+def h(x: int, y: str):
+    if x > 10:
+        return x
+    elif x > 5:
+        return y
+
+reveal_type(h(1, "a"))  # revealed: int | str | None
+
+def generator():
+    yield 1
+    yield 2
+    return None
+
+# TODO: Should be `Generator[Literal[1, 2], Any, None]`
+reveal_type(generator())  # revealed: None
+```
+
 ## Invalid return type
 
 <!-- snapshot-diagnostics -->

crates/ty_python_semantic/src/types.rs

@@ -4179,6 +4179,14 @@ impl<'db> Type<'db> {
         }
     }
 
+    /// Returns the inferred return type of `self` if it is a function literal.
+    fn inferred_return_type(self, db: &'db dyn Db) -> Option<Type<'db>> {
+        match self {
+            Type::FunctionLiteral(function_type) => Some(function_type.inferred_return_type(db)),
+            _ => None,
+        }
+    }
+
     /// Calls `self`. Returns a [`CallError`] if `self` is (always or possibly) not callable, or if
     /// the arguments are not compatible with the formal parameters.
     ///
@@ -4191,7 +4199,9 @@ impl<'db> Type<'db> {
         argument_types: &CallArgumentTypes<'_, 'db>,
     ) -> Result<Bindings<'db>, CallError<'db>> {
         let signatures = self.signatures(db);
-        Bindings::match_parameters(signatures, argument_types).check_types(db, argument_types)
+        let inferred_return_ty = || self.inferred_return_type(db).unwrap_or(Type::unknown());
+        Bindings::match_parameters(signatures, inferred_return_ty, argument_types)
+            .check_types(db, argument_types)
     }
 
     /// Look up a dunder method on the meta-type of `self` and call it.
@@ -4229,8 +4239,14 @@ impl<'db> Type<'db> {
         {
             Symbol::Type(dunder_callable, boundness) => {
                 let signatures = dunder_callable.signatures(db);
-                let bindings = Bindings::match_parameters(signatures, argument_types)
-                    .check_types(db, argument_types)?;
+                let inferred_return_ty = || {
+                    dunder_callable
+                        .inferred_return_type(db)
+                        .unwrap_or(Type::unknown())
+                };
+                let bindings =
+                    Bindings::match_parameters(signatures, inferred_return_ty, argument_types)
+                        .check_types(db, argument_types)?;
                 if boundness == Boundness::PossiblyUnbound {
                     return Err(CallDunderError::PossiblyUnbound(Box::new(bindings)));
                 }
@@ -6689,6 +6705,13 @@ impl<'db> FunctionType<'db> {
         signature
     }
 
+    /// Infers this function scope's types and returns the inferred return type.
+    fn inferred_return_type(self, db: &'db dyn Db) -> Type<'db> {
+        let scope = self.body_scope(db);
+        let inference = infer_scope_types(db, scope);
+        inference.inferred_return_type(db)
+    }
+
     pub(crate) fn is_known(self, db: &'db dyn Db, known_function: KnownFunction) -> bool {
         self.known(db) == Some(known_function)
     }

crates/ty_python_semantic/src/types/call/bind.rs

@@ -56,6 +56,7 @@ impl<'db> Bindings<'db> {
     /// verify that each argument type is assignable to the corresponding parameter type.
     pub(crate) fn match_parameters(
         signatures: Signatures<'db>,
+        inferred_return_ty: impl Fn() -> Type<'db> + Copy,
         arguments: &CallArguments<'_>,
     ) -> Self {
         let mut argument_forms = vec![None; arguments.len()];
@@ -68,6 +69,7 @@ impl<'db> Bindings<'db> {
                     arguments,
                     &mut argument_forms,
                     &mut conflicting_forms,
+                    inferred_return_ty,
                 )
             })
             .collect();
@@ -892,6 +894,7 @@ impl<'db> CallableBinding<'db> {
         arguments: &CallArguments<'_>,
         argument_forms: &mut [Option<ParameterForm>],
         conflicting_forms: &mut [bool],
+        inferred_return_ty: impl Fn() -> Type<'db> + Copy,
     ) -> Self {
         // If this callable is a bound method, prepend the self instance onto the arguments list
         // before checking.
@@ -911,6 +914,7 @@ impl<'db> CallableBinding<'db> {
                     arguments.as_ref(),
                     argument_forms,
                     conflicting_forms,
+                    inferred_return_ty,
                 )
             })
             .collect();
@@ -1079,6 +1083,7 @@ impl<'db> Binding<'db> {
         arguments: &CallArguments<'_>,
         argument_forms: &mut [Option<ParameterForm>],
         conflicting_forms: &mut [bool],
+        inferred_return_ty: impl Fn() -> Type<'db>,
     ) -> Self {
         let parameters = signature.parameters();
         // The parameter that each argument is matched with.
@@ -1195,7 +1200,7 @@ impl<'db> Binding<'db> {
         }
 
         Self {
-            return_ty: signature.return_ty.unwrap_or(Type::unknown()),
+            return_ty: signature.return_ty.unwrap_or_else(inferred_return_ty),
             specialization: None,
             inherited_specialization: None,
             argument_parameters: argument_parameters.into_boxed_slice(),

crates/ty_python_semantic/src/types/infer.rs

@@ -57,7 +57,7 @@ use crate::semantic_index::expression::{Expression, ExpressionKind};
 use crate::semantic_index::symbol::{
     FileScopeId, NodeWithScopeKind, NodeWithScopeRef, ScopeId, ScopeKind,
 };
-use crate::semantic_index::{semantic_index, EagerBindingsResult, SemanticIndex};
+use crate::semantic_index::{semantic_index, use_def_map, EagerBindingsResult, SemanticIndex};
 use crate::symbol::{
     builtins_module_scope, builtins_symbol, explicit_global_symbol,
     module_type_implicit_global_symbol, symbol, symbol_from_bindings, symbol_from_declarations,
@@ -366,6 +366,11 @@ pub(crate) struct TypeInference<'db> {
     /// The scope this region is part of.
     scope: ScopeId<'db>,
 
+    /// The returned types of this region (if this is a function body).
+    ///
+    /// These are stored in `Vec` to delay the creation of the union type as long as possible.
+    return_types: Vec<Type<'db>>,
+
     /// The fallback type for missing expressions/bindings/declarations.
     ///
     /// This is used only when constructing a cycle-recovery `TypeInference`.
@@ -381,6 +386,7 @@ impl<'db> TypeInference<'db> {
             deferred: FxHashSet::default(),
             diagnostics: TypeCheckDiagnostics::default(),
             scope,
+            return_types: vec![],
             cycle_fallback_type: None,
         }
     }
@@ -393,6 +399,7 @@ impl<'db> TypeInference<'db> {
             deferred: FxHashSet::default(),
             diagnostics: TypeCheckDiagnostics::default(),
             scope,
+            return_types: vec![],
             cycle_fallback_type: Some(cycle_fallback_type),
         }
     }
@@ -440,12 +447,27 @@ impl<'db> TypeInference<'db> {
         &self.diagnostics
     }
 
+    /// Returns the inferred return type of this function body (union of all possible return types),
+    /// or `None` if the region is not a function body.
+    pub(crate) fn inferred_return_type(&self, db: &'db dyn Db) -> Type<'db> {
+        let mut union = UnionBuilder::new(db);
+        for ty in &self.return_types {
+            union = union.add(*ty);
+        }
+        let use_def = use_def_map(db, self.scope);
+        if use_def.can_implicit_return(db) {
+            union = union.add(Type::none(db));
+        }
+        union.build()
+    }
+
     fn shrink_to_fit(&mut self) {
         self.expressions.shrink_to_fit();
         self.bindings.shrink_to_fit();
         self.declarations.shrink_to_fit();
         self.diagnostics.shrink_to_fit();
         self.deferred.shrink_to_fit();
+        self.return_types.shrink_to_fit();
     }
 }
 
@@ -4674,7 +4696,12 @@ impl<'db> TypeInferenceBuilder<'db> {
         }
 
         let signatures = callable_type.signatures(self.db());
-        let bindings = Bindings::match_parameters(signatures, &call_arguments);
+        let inferred_return_ty = || {
+            callable_type
+                .inferred_return_type(self.db())
+                .unwrap_or(Type::unknown())
+        };
+        let bindings = Bindings::match_parameters(signatures, inferred_return_ty, &call_arguments);
         let call_argument_types =
             self.infer_argument_types(arguments, call_arguments, &bindings.argument_forms);
 
@@ -6689,15 +6716,21 @@ impl<'db> TypeInferenceBuilder<'db> {
             value_ty,
             generic_context.signature(self.db()),
         ));
-        let bindings = match Bindings::match_parameters(signatures, &call_argument_types)
-            .check_types(self.db(), &call_argument_types)
-        {
-            Ok(bindings) => bindings,
-            Err(CallError(_, bindings)) => {
-                bindings.report_diagnostics(&self.context, subscript.into());
-                return Type::unknown();
-            }
+        let inferred_return_ty = || {
+            value_ty
+                .inferred_return_type(self.db())
+                .unwrap_or(Type::unknown())
         };
+        let bindings =
+            match Bindings::match_parameters(signatures, inferred_return_ty, &call_argument_types)
+                .check_types(self.db(), &call_argument_types)
+            {
+                Ok(bindings) => bindings,
+                Err(CallError(_, bindings)) => {
+                    bindings.report_diagnostics(&self.context, subscript.into());
+                    return Type::unknown();
+                }
+            };
         let callable = bindings
             .into_iter()
             .next()
@@ -7105,6 +7138,11 @@ impl<'db> TypeInferenceBuilder<'db> {
         self.infer_region();
         self.types.diagnostics = self.context.finish();
         self.types.shrink_to_fit();
+        self.types.return_types = self
+            .return_types_and_ranges
+            .into_iter()
+            .map(|ty_range| ty_range.ty)
+            .collect();
         self.types
     }
 }