Support lambdas

mypy/checker.py

@@ -4280,12 +4280,14 @@ def check_return_stmt(self, s: ReturnStmt) -> None:
                 return_type = self.return_types[-1]
             return_type = get_proper_type(return_type)
 
+            is_lambda = isinstance(self.scope.top_function(), LambdaExpr)
             if isinstance(return_type, UninhabitedType):
-                self.fail(message_registry.NO_RETURN_EXPECTED, s)
-                return
+                # Avoid extra error messages for failed inference in lambdas
+                if not is_lambda or not return_type.ambiguous:
+                    self.fail(message_registry.NO_RETURN_EXPECTED, s)
+                    return
 
             if s.expr:
-                is_lambda = isinstance(self.scope.top_function(), LambdaExpr)
                 declared_none_return = isinstance(return_type, NoneType)
                 declared_any_return = isinstance(return_type, AnyType)
 
@@ -7366,6 +7368,11 @@ def visit_erased_type(self, t: ErasedType) -> bool:
         # This can happen inside a lambda.
         return True
 
+    def visit_type_var(self, t: TypeVarType) -> bool:
+        # This is needed to prevent leaking into partial types during
+        # multi-step type inference.
+        return t.id.is_meta_var()
+
 
 class SetNothingToAny(TypeTranslator):
     """Replace all ambiguous <nothing> types with Any (to avoid spurious extra errors)."""

mypy/checkexpr.py

@@ -1858,6 +1858,8 @@ def infer_function_type_arguments_using_context(
             #        def identity(x: T) -> T: return x
             #
             #        expects_literal(identity(3))  # Should type-check
+            # TODO: we may want to add similar exception if all arguments are lambdas, since
+            # in this case external context is almost everything we have.
             if not is_generic_instance(ctx) and not is_literal_type_like(ctx):
                 return callable.copy_modified()
         args = infer_type_arguments(callable.variables, ret_type, erased_ctx)
@@ -4677,8 +4679,22 @@ def infer_lambda_type_using_context(
         # they must be considered as indeterminate. We use ErasedType since it
         # does not affect type inference results (it is for purposes like this
         # only).
-        callable_ctx = get_proper_type(replace_meta_vars(ctx, ErasedType()))
-        assert isinstance(callable_ctx, CallableType)
+        if self.chk.options.new_type_inference:
+            # With new type inference we can preserve argument types even if they
+            # are generic, since new inference algorithm can handle constraints
+            # like S <: T (we still erase return type since it's ultimately unknown).
+            extra_vars = []
+            for arg in ctx.arg_types:
+                meta_vars = [tv for tv in get_all_type_vars(arg) if tv.id.is_meta_var()]
+                extra_vars.extend([tv for tv in meta_vars if tv not in extra_vars])
+            callable_ctx = ctx.copy_modified(
+                ret_type=replace_meta_vars(ctx.ret_type, ErasedType()),
+                variables=list(ctx.variables) + extra_vars,
+            )
+        else:
+            erased_ctx = replace_meta_vars(ctx, ErasedType())
+            assert isinstance(erased_ctx, ProperType) and isinstance(erased_ctx, CallableType)
+            callable_ctx = erased_ctx
 
         # The callable_ctx may have a fallback of builtins.type if the context
         # is a constructor -- but this fallback doesn't make sense for lambdas.

test-data/unit/check-generics.test

@@ -2713,6 +2713,7 @@ reveal_type(func(1))  # N: Revealed type is "builtins.int"
 [builtins fixtures/tuple.pyi]
 
 [case testGenericLambdaGenericMethodNoCrash]
+# flags: --new-type-inference
 from typing import TypeVar, Union, Callable, Generic
 
 S = TypeVar("S")
@@ -2723,7 +2724,7 @@ def f(x: Callable[[G[T]], int]) -> T: ...
 class G(Generic[T]):
     def g(self, x: S) -> Union[S, T]: ...
 
-f(lambda x: x.g(0))  # E: Cannot infer type argument 1 of "f"
+f(lambda x: x.g(0))  # E: Incompatible return value type (got "Union[int, T]", expected "int")
 
 [case testDictStarInference]
 class B: ...
@@ -3036,6 +3037,34 @@ reveal_type(dec2(id1))  # N: Revealed type is "def [UC <: __main__.C] (UC`5) ->
 reveal_type(dec2(id2))  # N: Revealed type is "def (<nothing>) -> builtins.list[<nothing>]" \
                         # E: Argument 1 to "dec2" has incompatible type "Callable[[V], V]"; expected "Callable[[<nothing>], <nothing>]"
 
+[case testInferenceAgainstGenericLambdas]
+# flags: --new-type-inference
+from typing import TypeVar, Callable, List
+
+S = TypeVar('S')
+T = TypeVar('T')
+
+def dec1(f: Callable[[T], T]) -> Callable[[T], List[T]]:
+    ...
+def dec2(f: Callable[[S], T]) -> Callable[[S], List[T]]:
+    ...
+def dec3(f: Callable[[List[S]], T]) -> Callable[[S], T]:
+    ...
+def dec4(f: Callable[[S], List[T]]) -> Callable[[S], T]:
+    ...
+def dec5(f: Callable[[int], T]) -> Callable[[int], List[T]]:
+    def g(x: int) -> List[T]:
+        return [f(x)] * x
+    return g
+
+reveal_type(dec1(lambda x: x))  # N: Revealed type is "def [T] (T`2) -> builtins.list[T`2]"
+reveal_type(dec2(lambda x: x))  # N: Revealed type is "def [S] (S`3) -> builtins.list[S`3]"
+reveal_type(dec3(lambda x: x[0]))  # N: Revealed type is "def [S] (S`5) -> S`5"
+reveal_type(dec4(lambda x: [x]))  # N: Revealed type is "def [S] (S`7) -> S`7"
+reveal_type(dec1(lambda x: 1))  # N: Revealed type is "def (builtins.int) -> builtins.list[builtins.int]"
+reveal_type(dec5(lambda x: x))  # N: Revealed type is "def (builtins.int) -> builtins.list[builtins.int]"
+[builtins fixtures/list.pyi]
+
 [case testInferenceAgainstGenericParamSpecBasicInList]
 # flags: --new-type-inference
 from typing import TypeVar, Callable, List, Tuple

test-data/unit/check-inference-context.test

@@ -693,14 +693,15 @@ f(lambda: None)
 g(lambda: None)
 
 [case testIsinstanceInInferredLambda]
+# flags: --new-type-inference
 from typing import TypeVar, Callable, Optional
 T = TypeVar('T')
 S = TypeVar('S')
 class A: pass
 class B(A): pass
 class C(A): pass
 def f(func: Callable[[T], S], *z: T, r: Optional[S] = None) -> S: pass
-f(lambda x: 0 if isinstance(x, B) else 1) # E: Cannot infer type argument 1 of "f"
+reveal_type(f(lambda x: 0 if isinstance(x, B) else 1))  # N: Revealed type is "builtins.int"
 f(lambda x: 0 if isinstance(x, B) else 1, A())() # E: "int" not callable
 f(lambda x: x if isinstance(x, B) else B(), A(), r=B())() # E: "B" not callable
 f(

test-data/unit/check-inference.test

@@ -1375,19 +1375,21 @@ class B: pass
 [builtins fixtures/list.pyi]
 
 [case testUninferableLambda]
+# flags: --new-type-inference
 from typing import TypeVar, Callable
 X = TypeVar('X')
 def f(x: Callable[[X], X]) -> X: pass
-y = f(lambda x: x) # E: Cannot infer type argument 1 of "f"
+y = f(lambda x: x)  # E: Need type annotation for "y"
 
 [case testUninferableLambdaWithTypeError]
+# flags: --new-type-inference
 from typing import TypeVar, Callable
 X = TypeVar('X')
 def f(x: Callable[[X], X], y: str) -> X: pass
 y = f(lambda x: x, 1) # Fail
 [out]
-main:4: error: Cannot infer type argument 1 of "f"
-main:4: error: Argument 2 to "f" has incompatible type "int"; expected "str"
+main:5: error: Need type annotation for "y"
+main:5: error: Argument 2 to "f" has incompatible type "int"; expected "str"
 
 [case testInferLambdaNone]
 # flags: --no-strict-optional