[red-knot] Infer return type of lambda expression

Author: dhruvmanila

crates/red_knot_python_semantic/resources/mdtest/expression/lambda.md

@@ -5,10 +5,10 @@
 `lambda` expressions can be defined without any parameters.
 
 ```py
-reveal_type(lambda: 1)  # revealed: () -> @Todo(lambda return type)
+reveal_type(lambda: 1)  # revealed: () -> Literal[1]
 
 # error: [unresolved-reference]
-reveal_type(lambda: a)  # revealed: () -> @Todo(lambda return type)
+reveal_type(lambda: a)  # revealed: () -> Unknown
 ```
 
 ## With parameters
@@ -17,45 +17,47 @@ Unlike parameters in function definition, the parameters in a `lambda` expressio
 annotated.
 
 ```py
-reveal_type(lambda a: a)  # revealed: (a) -> @Todo(lambda return type)
-reveal_type(lambda a, b: a + b)  # revealed: (a, b) -> @Todo(lambda return type)
+reveal_type(lambda a: a)  # revealed: (a) -> Unknown
+reveal_type(lambda a, b: a + b)  # revealed: (a, b) -> Unknown
 ```
 
 But, it can have default values:
 
 ```py
-reveal_type(lambda a=1: a)  # revealed: (a=Literal[1]) -> @Todo(lambda return type)
-reveal_type(lambda a, b=2: a)  # revealed: (a, b=Literal[2]) -> @Todo(lambda return type)
+reveal_type(lambda a=1: a)  # revealed: (a=Literal[1]) -> Unknown | Literal[1]
+reveal_type(lambda a, b=2: a)  # revealed: (a, b=Literal[2]) -> Unknown
 ```
 
 And, positional-only parameters:
 
 ```py
-reveal_type(lambda a, b, /, c: c)  # revealed: (a, b, /, c) -> @Todo(lambda return type)
+reveal_type(lambda a, b, /, c: c)  # revealed: (a, b, /, c) -> Unknown
 ```
 
 And, keyword-only parameters:
 
 ```py
-reveal_type(lambda a, *, b=2, c: b)  # revealed: (a, *, b=Literal[2], c) -> @Todo(lambda return type)
+reveal_type(lambda a, *, b=2, c: b)  # revealed: (a, *, b=Literal[2], c) -> Unknown | Literal[2]
 ```
 
 And, variadic parameter:
 
 ```py
-reveal_type(lambda *args: args)  # revealed: (*args) -> @Todo(lambda return type)
+# TODO: should be `tuple[Unknown, ...]` (needs generics)
+reveal_type(lambda *args: args)  # revealed: (*args) -> tuple
 ```
 
 And, keyword-varidic parameter:
 
 ```py
-reveal_type(lambda **kwargs: kwargs)  # revealed: (**kwargs) -> @Todo(lambda return type)
+# TODO: should be `dict[str, Unknown]` (needs generics)
+reveal_type(lambda **kwargs: kwargs)  # revealed: (**kwargs) -> dict
 ```
 
 Mixing all of them together:
 
 ```py
-# revealed: (a, b, /, c=Literal[True], *args, *, d=Literal["default"], e=Literal[5], **kwargs) -> @Todo(lambda return type)
+# revealed: (a, b, /, c=Literal[True], *args, *, d=Literal["default"], e=Literal[5], **kwargs) -> None
 reveal_type(lambda a, b, /, c=True, *args, d="default", e=5, **kwargs: None)
 ```
 
@@ -96,5 +98,5 @@ Here, a `lambda` expression is used as the default value for a parameter in anot
 expression.
 
 ```py
-reveal_type(lambda a=lambda x, y: 0: 2)  # revealed: (a=(x, y) -> @Todo(lambda return type)) -> @Todo(lambda return type)
+reveal_type(lambda a=lambda x, y: 0: 2)  # revealed: (a=(x, y) -> Literal[0]) -> Literal[2]
 ```

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -913,6 +913,7 @@ where
                     .iter_non_variadic_params()
                     .filter_map(|param| param.default.as_deref())
                 {
+                    self.add_standalone_expression(default);
                     self.visit_expr(default);
                 }
                 // The symbol for the function name itself has to be evaluated
@@ -1738,6 +1739,7 @@ where
                         .iter_non_variadic_params()
                         .filter_map(|param| param.default.as_deref())
                     {
+                        self.add_standalone_expression(default);
                         self.visit_expr(default);
                     }
                     self.visit_parameters(parameters);

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1338,7 +1338,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             .iter_non_variadic_params()
             .filter_map(|param| param.default.as_deref())
         {
-            self.infer_expression(default);
+            self.infer_standalone_expression(default);
         }
 
         // If there are type params, parameters and returns are evaluated in that scope, that is, in
@@ -1463,7 +1463,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         } = parameter_with_default;
         let default_ty = default
             .as_ref()
-            .map(|default| self.file_expression_type(default));
+            .map(|default| infer_expression_type(self.db(), self.index.expression(&**default)));
         if let Some(annotation) = parameter.annotation.as_ref() {
             let declared_ty = self.file_expression_type(annotation);
             let declared_and_inferred_ty = if let Some(default_ty) = default_ty {
@@ -3447,9 +3447,15 @@ impl<'db> TypeInferenceBuilder<'db> {
         let ast::ExprLambda {
             range: _,
             parameters,
-            body: _,
+            body,
         } = lambda_expression;
 
+        let mut default_type = |parameter: &ast::ParameterWithDefault| {
+            parameter
+                .default()
+                .map(|default| self.infer_standalone_expression(default))
+        };
+
         let parameters = if let Some(parameters) = parameters {
             let positional_only = parameters
                 .posonlyargs
@@ -3459,9 +3465,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                         Some(parameter.name().id.clone()),
                         None,
                         ParameterKind::PositionalOnly {
-                            default_ty: parameter
-                                .default()
-                                .map(|default| self.infer_expression(default)),
+                            default_ty: default_type(parameter),
                         },
                     )
                 })
@@ -3474,9 +3478,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                         Some(parameter.name().id.clone()),
                         None,
                         ParameterKind::PositionalOrKeyword {
-                            default_ty: parameter
-                                .default()
-                                .map(|default| self.infer_expression(default)),
+                            default_ty: default_type(parameter),
                         },
                     )
                 })
@@ -3496,9 +3498,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                         Some(parameter.name().id.clone()),
                         None,
                         ParameterKind::KeywordOnly {
-                            default_ty: parameter
-                                .default()
-                                .map(|default| self.infer_expression(default)),
+                            default_ty: default_type(parameter),
                         },
                     )
                 })
@@ -3525,7 +3525,7 @@ impl<'db> TypeInferenceBuilder<'db> {
 
         Type::Callable(CallableType::General(GeneralCallableType::new(
             self.db(),
-            Signature::new(parameters, Some(todo_type!("lambda return type"))),
+            Signature::new(parameters, Some(self.file_expression_type(body))),
         )))
     }