From b7d8724f5028fff5e6c1bfe86334bcdec79177dc Mon Sep 17 00:00:00 2001
From: Nikita Sobolev <mail@sobolevn.me>
Date: Wed, 22 Sep 2021 13:36:02 +0300
Subject: [PATCH] Adds `in_checked_function()` to semanal similar to the one in
 checker (#11156)

Refs #11155
---
 mypy/checkexpr.py |  3 +++
 mypy/semanal.py   | 21 ++++++++++++++-------
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index 30db1e8ac87d..d32e9f8eafb5 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -3506,6 +3506,9 @@ def visit_lambda_expr(self, e: LambdaExpr) -> Type:
             # Type check everything in the body except for the final return
             # statement (it can contain tuple unpacking before return).
             with self.chk.scope.push_function(e):
+                # Lambdas can have more than one element in body,
+                # when we add "fictional" AssigmentStatement nodes, like in:
+                # `lambda (a, b): a`
                 for stmt in e.body.body[:-1]:
                     stmt.accept(self.chk)
                 # Only type check the return expression, not the return statement.
diff --git a/mypy/semanal.py b/mypy/semanal.py
index 3f0c84e5a7c2..c59c48201dd0 100644
--- a/mypy/semanal.py
+++ b/mypy/semanal.py
@@ -4953,6 +4953,18 @@ def is_local_name(self, name: str) -> bool:
         """Does name look like reference to a definition in the current module?"""
         return self.is_defined_in_current_module(name) or '.' not in name
 
+    def in_checked_function(self) -> bool:
+        """Should we type-check the current function?
+
+        - Yes if --check-untyped-defs is set.
+        - Yes outside functions.
+        - Yes in annotated functions.
+        - No otherwise.
+        """
+        return (self.options.check_untyped_defs
+                or not self.function_stack
+                or not self.function_stack[-1].is_dynamic())
+
     def fail(self,
              msg: str,
              ctx: Context,
@@ -4960,10 +4972,7 @@ def fail(self,
              *,
              code: Optional[ErrorCode] = None,
              blocker: bool = False) -> None:
-        if (not serious and
-                not self.options.check_untyped_defs and
-                self.function_stack and
-                self.function_stack[-1].is_dynamic()):
+        if not serious and not self.in_checked_function():
             return
         # In case it's a bug and we don't really have context
         assert ctx is not None, msg
@@ -4973,9 +4982,7 @@ def fail_blocker(self, msg: str, ctx: Context) -> None:
         self.fail(msg, ctx, blocker=True)
 
     def note(self, msg: str, ctx: Context, code: Optional[ErrorCode] = None) -> None:
-        if (not self.options.check_untyped_defs and
-                self.function_stack and
-                self.function_stack[-1].is_dynamic()):
+        if not self.in_checked_function():
             return
         self.errors.report(ctx.get_line(), ctx.get_column(), msg, severity='note', code=code)