Support None partial types with local partial types (python#20938)

This turned out to be surprisingly easy, hopefully I didn't miss
anything. Implementation is a bit ad-hoc, but the feature is equally
ad-hoc. I enable both class-level and module-level `None` cases.

Note this also naturally fixes a bug when skipping function bodies in
third-party packages. I added this as an optimization of the previous
stripping bodies optimization, but I forgot about partial types.

Author: ilevkivskyi

mypy/checker.py

@@ -719,10 +719,25 @@ def accept_loop(
     # Definitions
     #
 
+    @contextmanager
+    def set_recurse_into_functions(self) -> Iterator[None]:
+        """Temporarily set recurse_into_functions to True.
+
+        This is used to process top-level functions/methods as a whole.
+        """
+        old_recurse_into_functions = self.recurse_into_functions
+        self.recurse_into_functions = True
+        try:
+            yield
+        finally:
+            self.recurse_into_functions = old_recurse_into_functions
+
     def visit_overloaded_func_def(self, defn: OverloadedFuncDef) -> None:
-        if not self.recurse_into_functions:
+        # If a function/method can infer variable types, it should be processed as part
+        # of the module top level (i.e. module interface).
+        if not self.recurse_into_functions and not defn.can_infer_vars:
             return
-        with self.tscope.function_scope(defn):
+        with self.tscope.function_scope(defn), self.set_recurse_into_functions():
             self._visit_overloaded_func_def(defn)
 
     def _visit_overloaded_func_def(self, defn: OverloadedFuncDef) -> None:
@@ -1196,9 +1211,9 @@ def get_generator_return_type(self, return_type: Type, is_coroutine: bool) -> Ty
             return NoneType()
 
     def visit_func_def(self, defn: FuncDef) -> None:
-        if not self.recurse_into_functions:
+        if not self.recurse_into_functions and not defn.can_infer_vars:
             return
-        with self.tscope.function_scope(defn):
+        with self.tscope.function_scope(defn), self.set_recurse_into_functions():
             self.check_func_item(defn, name=defn.name)
             if not self.can_skip_diagnostics:
                 if defn.info:
@@ -1438,6 +1453,7 @@ def check_func_def(
                         or self.options.preserve_asts
                         or not isinstance(defn, FuncDef)
                         or defn.has_self_attr_def
+                        or defn.can_infer_vars
                     ):
                         self.accept(item.body)
                 unreachable = self.binder.is_unreachable()
@@ -5604,8 +5620,8 @@ def visit_decorator(self, e: Decorator) -> None:
     def visit_decorator_inner(
         self, e: Decorator, allow_empty: bool = False, skip_first_item: bool = False
     ) -> None:
-        if self.recurse_into_functions:
-            with self.tscope.function_scope(e.func):
+        if self.recurse_into_functions or e.func.can_infer_vars:
+            with self.tscope.function_scope(e.func), self.set_recurse_into_functions():
                 self.check_func_item(e.func, name=e.func.name, allow_empty=allow_empty)
 
         # Process decorators from the inside out to determine decorated signature, which
@@ -7718,30 +7734,7 @@ def enter_partial_types(
         partial_types, _, _ = self.partial_types.pop()
         if not self.current_node_deferred:
             for var, context in partial_types.items():
-                # If we require local partial types, there are a few exceptions where
-                # we fall back to inferring just "None" as the type from a None initializer:
-                #
-                # 1. If all happens within a single function this is acceptable, since only
-                #    the topmost function is a separate target in fine-grained incremental mode.
-                #    We primarily want to avoid "splitting" partial types across targets.
-                #
-                # 2. A None initializer in the class body if the attribute is defined in a base
-                #    class is fine, since the attribute is already defined and it's currently okay
-                #    to vary the type of an attribute covariantly. The None type will still be
-                #    checked for compatibility with base classes elsewhere. Without this exception
-                #    mypy could require an annotation for an attribute that already has been
-                #    declared in a base class, which would be bad.
-                allow_none = (
-                    not self.options.local_partial_types
-                    or is_function
-                    or (is_class and self.is_defined_in_base_class(var))
-                )
-                if (
-                    allow_none
-                    and isinstance(var.type, PartialType)
-                    and var.type.type is None
-                    and not permissive
-                ):
+                if isinstance(var.type, PartialType) and var.type.type is None and not permissive:
                     var.type = NoneType()
                 else:
                     if var not in self.partial_reported and not permissive:
@@ -7812,10 +7805,15 @@ def find_partial_types_in_all_scopes(
                 # for fine-grained incremental mode).
                 disallow_other_scopes = self.options.local_partial_types
 
+                # There are two exceptions:
                 if isinstance(var.type, PartialType) and var.type.type is not None and var.info:
-                    # This is an ugly hack to make partial generic self attributes behave
-                    # as if --local-partial-types is always on (because it used to be like this).
+                    # We always prohibit non-None partial types at class scope
+                    # for historical reasons.
                     disallow_other_scopes = True
+                if isinstance(var.type, PartialType) and var.type.type is None:
+                    # We always allow None partial types, since this is a common use case.
+                    # It is special-cased in fine-grained incremental mode.
+                    disallow_other_scopes = False
 
                 scope_active = (
                     not disallow_other_scopes or scope.is_local == self.partial_types[-1].is_local

mypy/nodes.py

@@ -589,6 +589,7 @@ class FuncBase(Node):
         "is_final",  # Uses "@final"
         "is_explicit_override",  # Uses "@override"
         "is_type_check_only",  # Uses "@type_check_only"
+        "can_infer_vars",
         "_fullname",
     )
 
@@ -609,6 +610,18 @@ def __init__(self) -> None:
         self.is_final = False
         self.is_explicit_override = False
         self.is_type_check_only = False
+        # Can this function/method infer types of variables defined outside? Currently,
+        # we only set this in cases like:
+        #     x = None
+        #     def foo() -> None:
+        #         global x
+        #         x = 1
+        # and
+        #     class C:
+        #         x = None
+        #         def foo(self) -> None:
+        #             self.x = 1
+        self.can_infer_vars = False
         # Name with module prefix
         self._fullname = ""
 

mypy/semanal.py

@@ -4589,6 +4589,17 @@ def make_name_lvalue_point_to_existing_def(
                 self.name_not_defined(lval.name, lval)
             self.check_lvalue_validity(lval.node, lval)
 
+        if self.scope.functions and lval.name in self.global_decls[-1]:
+            # Technically, we only need to set this if original r.h.s. would be inferred
+            # as None, but it is tricky to detect reliably during semantic analysis.
+            if (
+                original_def
+                and isinstance(original_def.node, Var)
+                and original_def.node.is_inferred
+            ):
+                for func in self.scope.functions:
+                    func.can_infer_vars = True
+
     def analyze_tuple_or_list_lvalue(self, lval: TupleExpr, explicit_type: bool = False) -> None:
         """Analyze an lvalue or assignment target that is a list or tuple."""
         items = lval.items
@@ -4674,6 +4685,14 @@ def analyze_member_lvalue(
                     for func in self.scope.functions:
                         if isinstance(func, FuncDef):
                             func.has_self_attr_def = True
+            if (
+                cur_node
+                and isinstance(cur_node.node, Var)
+                and cur_node.node.is_inferred
+                and cur_node.node.is_initialized_in_class
+            ):
+                for func in self.scope.functions:
+                    func.can_infer_vars = True
         self.check_lvalue_validity(lval.node, lval)
 
     def is_self_member_ref(self, memberexpr: MemberExpr) -> bool:

mypy/server/update.py

@@ -138,6 +138,7 @@
 from mypy.modulefinder import BuildSource
 from mypy.nodes import (
     Decorator,
+    FuncBase,
     FuncDef,
     ImportFrom,
     MypyFile,
@@ -953,6 +954,20 @@ def find_targets_recursive(
                 deferred, stale_proto = lookup_target(manager, target)
                 if stale_proto:
                     stale_protos.add(stale_proto)
+
+                # If there are function targets that can infer outer variables, they should
+                # be re-processed as part of the module top-level instead (for consistency).
+                regular = []
+                shared = []
+                for d in deferred:
+                    if isinstance(d.node, FuncBase) and d.node.can_infer_vars:
+                        shared.append(d)
+                    else:
+                        regular.append(d)
+                deferred = regular
+                if shared:
+                    deferred.append(FineGrainedDeferredNode(manager.modules[module_id], None))
+
                 result[module_id].update(deferred)
 
     return result, unloaded_files, stale_protos