Elaboration

mypy/build.py

@@ -212,6 +212,8 @@
 # situations where 100 empty __init__.py files cost less than 1 trivial module.
 MIN_SIZE_HINT: Final = 256
 
+MAX_PLUGIN_REFINEMENT_PASSES = 5
+
 
 class SCC:
     """A simple class that represents a strongly connected component (import cycle)."""
@@ -3964,7 +3966,7 @@ def module_not_found(
                     errors.report(
                         line,
                         0,
-                        f'Did you mean {pretty_seq(matches, "or")}?',
+                        f"Did you mean {pretty_seq(matches, 'or')}?",
                         severity="note",
                         code=code,
                     )
@@ -4209,7 +4211,7 @@ def dump_line_checking_stats(path: str, graph: Graph) -> None:
             f.write(f"{id}:\n")
             for line in sorted(graph[id].per_line_checking_time_ns):
                 line_time = graph[id].per_line_checking_time_ns[line]
-                f.write(f"{line:>5} {line_time/1000:8.1f}\n")
+                f.write(f"{line:>5} {line_time / 1000:8.1f}\n")
 
 
 def dump_graph(graph: Graph, stdout: TextIO | None = None) -> None:
@@ -4786,27 +4788,70 @@ def process_stale_scc(graph: Graph, ascc: SCC, manager: BuildManager) -> None:
     mypy.semanal_main.semantic_analysis_for_scc(graph, scc, manager.errors)
 
     t3 = time.time()
-    # Track what modules aren't yet done, so we can finish them as soon
-    # as possible, saving memory.
-    unfinished_modules = set(stale)
-    for id in stale:
-        graph[id].type_check_first_pass()
-        if not graph[id].type_checker().deferred_nodes:
-            unfinished_modules.discard(id)
-            graph[id].detect_possibly_undefined_vars()
-            graph[id].finish_passes()
+    # A plugin might choose to refine callable return types after a checking pass.
+    # In any case this happens, we want to recheck scc until fixed point, and avoid
+    # broadcasting errors that violate the domain specific typing rule.
+    pre_type_check_errors = {
+        id: list(manager.errors.error_info_map.get(graph[id].xpath, [])) for id in stale
+    }
+    refinement_round = 0
+    while True:
+        # Track what modules aren't yet done, so we can finish them as soon
+        # as possible, saving memory.
+        unfinished_modules = set(stale)
 
-    while unfinished_modules:
         for id in stale:
-            if id not in unfinished_modules:
-                continue
-            if not graph[id].type_check_second_pass():
+            graph[id].type_check_first_pass()
+            if not graph[id].type_checker().deferred_nodes:
                 unfinished_modules.discard(id)
-                graph[id].detect_possibly_undefined_vars()
-                graph[id].finish_passes()
+
+        while unfinished_modules:
+            for id in stale:
+                if id not in unfinished_modules:
+                    continue
+                if not graph[id].type_check_second_pass():
+                    unfinished_modules.discard(id)
+
+        dirty_signatures = set()
+        for id in stale:
+            dirty_signatures |= graph[id].type_checker().take_changed_plugin_signatures()
+        if not dirty_signatures:
+            break
+
+        if refinement_round >= MAX_PLUGIN_REFINEMENT_PASSES:
+            raise RuntimeError(
+                "Function-body plugin refinements did not converge in SCC: "
+                + ", ".join(sorted(stale))
+            )
+        refinement_round += 1
+
+        for id in stale:
+            state = graph[id]
+            path = state.xpath
+
+            targets = {
+                info.target
+                for info in manager.errors.error_info_map.get(path, [])
+                if info.target is not None
+            }
+            manager.errors.clear_errors_in_targets(path, targets)
+
+            # Restore diagnostics emitted before type checking, such as semantic
+            # analysis errors. clear_errors_in_targets() removed these too when
+            # they belonged to a target we rechecked.
+            for info in pre_type_check_errors[id]:
+                if info.target in targets:
+                    manager.errors.add_error_info(info, file=path)
+
+            checker = graph[id].type_checker()
+            checker.reset()
+            checker.pass_num = 0
+
     for id in stale:
+        graph[id].detect_possibly_undefined_vars()
         graph[id].generate_unused_ignore_notes()
         graph[id].generate_ignore_without_code_notes()
+        graph[id].finish_passes()
 
     t4 = time.time()
     # Flush errors, and write cache in two phases: first data files, then meta files.

mypy/checker.py

@@ -212,7 +212,13 @@ def __init__(self) -> None:
 from mypy.operators import flip_ops, int_op_to_method, neg_ops
 from mypy.options import PRECISE_TUPLE_TYPES, Options
 from mypy.patterns import AsPattern, StarredPattern
-from mypy.plugin import Plugin
+from mypy.plugin import (
+    FunctionBodyContext,
+    FunctionBodyHook,
+    FunctionDefContext,
+    Plugin,
+    ReturnSite,
+)
 from mypy.plugins import dataclasses as dataclasses_plugin
 from mypy.scope import Scope
 from mypy.semanal import is_trivial_body, refers_to_fullname, set_callable_name
@@ -232,7 +238,12 @@ def __init__(self) -> None:
     restrict_subtype_away,
     unify_generic_callable,
 )
-from mypy.traverser import TraverserVisitor, all_return_statements, has_return_statement
+from mypy.traverser import (
+    TraverserVisitor,
+    all_return_statements,
+    can_fall_through,
+    has_return_statement,
+)
 from mypy.treetransform import TransformVisitor
 from mypy.typeanal import check_for_explicit_any, has_any_from_unimported_type, make_optional_type
 from mypy.typeops import (
@@ -334,6 +345,14 @@ class FineGrainedDeferredNode(NamedTuple):
     active_typeinfo: TypeInfo | None
 
 
+# Solely a plugin extension that tracks extra info internally while mypy visits a function body
+class ActiveFunctionHook(NamedTuple):
+    defn: FuncDef
+    declared_signature: CallableType
+    callback: FunctionBodyHook
+    return_sites: list[ReturnSite]
+
+
 # Data structure returned by find_isinstance_check representing
 # information learned from the truth or falsehood of a condition.  The
 # dict maps nodes representing expressions like 'a[0].x' to their
@@ -461,7 +480,10 @@ class TypeChecker(NodeVisitor[None], TypeCheckerSharedApi, SplittingVisitor):
 
     # Plugin that provides special type checking rules for specific library
     # functions such as open(), etc.
+    # and private attributes to support elaboration
     plugin: Plugin
+    _function_body_hooks: list[ActiveFunctionHook]
+    _changed_plugin_signatures: set[str]
 
     # A helper state to produce unique temporary names on demand.
     _unique_id: int
@@ -512,6 +534,8 @@ def __init__(
         self.inferred_attribute_types = None
         self.allow_constructor_cache = True
         self.local_type_map = LocalTypeMap(self)
+        self._function_body_hooks = []
+        self._changed_plugin_signatures = set()
 
         self.can_skip_diagnostics: Final = (
             self.options.ignore_errors
@@ -1333,6 +1357,77 @@ def get_generator_return_type(self, return_type: Type, is_coroutine: bool) -> Ty
             # Treat `Iterator[X]` as a shorthand for `Generator[X, Any, None]`.
             return NoneType()
 
+    def record_plugin_return(self, statement: ReturnStmt, inferred_type: Type) -> None:
+        if not self._function_body_hooks:
+            return
+
+        active = self._function_body_hooks[-1]
+
+        if self.scope.current_function() is not active.defn:
+            return
+
+        active.return_sites.append(ReturnSite(statement=statement, inferred_type=inferred_type))
+
+    def publish_plugin_refinement(
+        self, defn: FuncDef, declared: CallableType, refined_return: Type
+    ) -> None:
+        new_signature = declared.copy_modified(ret_type=refined_return)
+
+        old_signature = defn.plugin_effective_type or declared
+
+        if is_same_type(old_signature, new_signature):
+            return
+
+        defn.plugin_effective_type = new_signature
+        self._changed_plugin_signatures.add(defn.fullname)
+
+    def finish_function_body_hook(self, active: ActiveFunctionHook) -> None:
+        inferred_typ = make_simplified_union([site.inferred_type for site in active.return_sites])
+
+        can_fall = can_fall_through(active.defn.body)
+        if can_fall:
+            inferred_typ = make_simplified_union([inferred_typ, NoneType()])
+
+        result = active.callback(
+            FunctionBodyContext(
+                definition=active.defn,
+                declared_signature=active.declared_signature,
+                inferred_return_type=inferred_typ,
+                return_sites=tuple(active.return_sites),
+                can_fall_through=can_fall,
+                api=self,
+            )
+        )
+
+        if result is not None and result.refined_return is not None:
+            self.publish_plugin_refinement(
+                active.defn, active.declared_signature, result.refined_return
+            )
+
+    def take_changed_plugin_signatures(self) -> set[str]:
+        """Return and clear public signatures changed by function-body plugins."""
+        changed = self._changed_plugin_signatures
+        self._changed_plugin_signatures = set()
+        return changed
+
+    @contextmanager
+    def function_def_hook(self, defn: FuncDef) -> Iterator[None]:
+        active: ActiveFunctionHook | None = None
+        hook = self.plugin.get_function_def_hook(defn.fullname)
+        if hook is not None and isinstance(defn.type, CallableType):
+            hook_result = hook(FunctionDefContext(defn, defn.type, self))
+            if hook_result is not None and hook_result.after_body is not None:
+                active = ActiveFunctionHook(defn, defn.type, hook_result.after_body, [])
+                self._function_body_hooks.append(active)
+        try:
+            yield
+        finally:
+            if active is not None:
+                assert self._function_body_hooks[-1] is active
+                self._function_body_hooks.pop()
+                if not self.current_node_deferred:
+                    self.finish_function_body_hook(active)
+
     def visit_func_def(self, defn: FuncDef) -> None:
         # Type check initialization expressions as part of top-level.
         if not self.can_skip_diagnostics:
@@ -1346,7 +1441,11 @@ def visit_func_def(self, defn: FuncDef) -> None:
         self.visit_func_def_impl(defn)
 
     def visit_func_def_impl(self, defn: FuncDef) -> None:
-        with self.tscope.function_scope(defn), self.set_recurse_into_functions():
+        with (
+            self.tscope.function_scope(defn),
+            self.set_recurse_into_functions(),
+            self.function_def_hook(defn),
+        ):
             self.check_func_item(defn, name=defn.name)
             if not self.can_skip_diagnostics:
                 if defn.info:
@@ -5322,6 +5421,8 @@ def check_return_stmt(self, s: ReturnStmt) -> None:
                             s.expr, return_type, allow_none_return=allow_none_func_call
                         )
                     )
+
+                self.record_plugin_return(s, typ)
                 # Treat NotImplemented as having type Any, consistent with its
                 # definition in typeshed prior to python/typeshed#4222.
                 if isinstance(typ, Instance) and typ.type.fullname in NOT_IMPLEMENTED_TYPE_NAMES:
@@ -8184,6 +8285,9 @@ def iterable_item_type(self, it: ProperType, context: Context) -> Type:
         return self.analyze_iterable_item_type_without_expression(it, context)[1]
 
     def function_type(self, func: FuncBase) -> FunctionLike:
+        if func.plugin_effective_type is not None:
+            return func.plugin_effective_type
+
         typ = function_type(func, self.named_type("builtins.function"))
         if (
             isinstance(func, FuncItem)

mypy/nodes.py

@@ -720,6 +720,7 @@ class FuncBase(Node):
     __slots__ = (
         "type",
         "unanalyzed_type",
+        "plugin_effective_type",
         "info",
         "is_property",
         "is_class",  # Uses "@classmethod" (explicit or implicit)
@@ -738,6 +739,8 @@ def __init__(self) -> None:
         self.type: mypy.types.ProperType | None = None
         # Original, not semantically analyzed type (used for reprocessing)
         self.unanalyzed_type: mypy.types.ProperType | None = None
+        # Inferred types from return sites (used for elaboration)
+        self.plugin_effective_type: mypy.types.FunctionLike | None = None
         # If method, reference to TypeInfo
         self.info = FUNC_NO_INFO
         self.is_property = False

mypy/plugin.py

@@ -134,7 +134,9 @@ class C: pass
     ClassDef,
     Context,
     Expression,
+    FuncDef,
     MypyFile,
+    ReturnStmt,
     SymbolTableNode,
     TypeInfo,
 )
@@ -517,6 +519,42 @@ class DynamicClassDefContext(NamedTuple):
     api: SemanticAnalyzerPluginInterface
 
 
+# A context for a function hook signature after available after decl.
+class FunctionDefContext(NamedTuple):
+    definition: FuncDef
+    declared_signature: CallableType
+    api: CheckerPluginInterface
+
+
+class ReturnSite(NamedTuple):
+    statement: ReturnStmt
+    inferred_type: Type
+
+
+# A context for a function hook after the body is checked.
+class FunctionBodyContext(NamedTuple):
+    definition: FuncDef
+    declared_signature: CallableType  # preserves the user annotations
+    inferred_return_type: Type  # inferred from return exprs in func, can be None
+    return_sites: tuple[ReturnSite, ...]
+    can_fall_through: bool  # tracks definitions with multiple exists (e.g early returns)
+    api: CheckerPluginInterface
+
+
+# None here means inferred function return type is not propagated
+class FunctionBodyResult(NamedTuple):
+    refined_return: Type | None = None
+
+
+# Supplying this callback tells mypy that this body's inferred information may affect the function's published signature.
+class FunctionDefHookResult(NamedTuple):
+    after_body: FunctionBodyHook | None = None
+
+
+FunctionBodyHook = Callable[[FunctionBodyContext], FunctionBodyResult | None]
+FunctionDefHook = Callable[[FunctionDefContext], FunctionDefHookResult | None]
+
+
 @mypyc_attr(allow_interpreted_subclasses=True)
 class Plugin(CommonPluginApi):
     """Base class of all type checker plugins.
@@ -819,6 +857,17 @@ def get_dynamic_class_hook(
         """
         return None
 
+    def get_function_def_hook(self, fullname: str) -> FunctionDefHook | None:
+        """Implement domain specific type checking logic of a function
+        or method definition.
+
+        The callback runs after its declared signature has been analyzed,
+        but before its body is finalized.
+
+        TODO: point to examples.
+        """
+        return None
+
 
 T = TypeVar("T")
 
@@ -884,6 +933,9 @@ def get_function_signature_hook(
     def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
         return self._find_hook(lambda plugin: plugin.get_function_hook(fullname))
 
+    def get_function_def_hook(self, fullname: str) -> FunctionDefHook | None:
+        return self._find_hook(lambda plugin: plugin.get_function_def_hook(fullname))
+
     def get_method_signature_hook(
         self, fullname: str
     ) -> Callable[[MethodSigContext], FunctionLike] | None: