diff --git a/mypy/build.py b/mypy/build.py index 96a67105c816c..6d0047df7f327 100644 --- a/mypy/build.py +++ b/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. diff --git a/mypy/checker.py b/mypy/checker.py index 33ed5387554d8..ec417c504a60f 100644 --- a/mypy/checker.py +++ b/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) diff --git a/mypy/nodes.py b/mypy/nodes.py index 53efc1fbc86b3..562e811844e6c 100644 --- a/mypy/nodes.py +++ b/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 diff --git a/mypy/plugin.py b/mypy/plugin.py index d0e5add71a5b5..cc36214cb0309 100644 --- a/mypy/plugin.py +++ b/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: diff --git a/mypy/traverser.py b/mypy/traverser.py index 224d3b3460eb4..a1cf828e45bae 100644 --- a/mypy/traverser.py +++ b/mypy/traverser.py @@ -964,6 +964,38 @@ def has_return_statement(fdef: FuncBase) -> bool: return seeker.found +class AlwaysExits(TraverserVisitor): + def __init__(self) -> None: + self.exists = False + + def always_exits(self, stmt: Node) -> bool: + if isinstance(stmt, Block): + return any(self.always_exits(s) for s in stmt.body) + if isinstance(stmt, (ReturnStmt | RaiseStmt)): + return True + if isinstance(stmt, IfStmt): + if stmt.else_body is None: + return False + return all(self.always_exits(b) for b in stmt.body) and self.always_exits( + stmt.else_body + ) + if isinstance(stmt, WithStmt): + return self.always_exits(stmt.body) + return False + + def visit_block(self, o: Block) -> None: + self.exists = self.always_exits(o) + + +def can_fall_through(body: Block) -> bool: + """Whether the function body may complete without an explicit return + or raise on some path, i.e. whether an implicit `return None` is + possible.""" + seeker = AlwaysExits() + body.accept(seeker) + return not seeker.exists + + class NameAndMemberCollector(TraverserVisitor): def __init__(self) -> None: super().__init__() diff --git a/test-data/unit/check-elaboration-plugin.test b/test-data/unit/check-elaboration-plugin.test new file mode 100644 index 0000000000000..fb7ed7e242b00 --- /dev/null +++ b/test-data/unit/check-elaboration-plugin.test @@ -0,0 +1,189 @@ +[case testElaboratesAnyReturn] +[builtins fixtures/elaboration.pyi] +# flags: --python-version 3.11 --config-file tmp/mypy.ini + +from typing import Any + +def bad(x: int) -> Any: + return str(x) + +reveal_type(bad(1)) # N: Revealed type is "builtins.str" +value: int = bad(1) # E: Incompatible types in assignment (expression has type "str", variable has type "int") + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py + +[case testUnionFromIssue] +# flags: --python-version 3.11 --config-file tmp/mypy.ini --show-error-codes +def some_method_with_union_return(some_input: str | int) -> str | int: # E: Union return types are not allowed [union-return] + return some_input + + +def some_other_method() -> None: + some_method_with_union_return(5) + + +def another_method() -> None: + some_method_with_union_return("foo") + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py + +[case testElaboratesNestedCallableReturn] +[builtins fixtures/tuple.pyi] +# flags: --python-version 3.11 --config-file tmp/mypy.ini --show-error-codes + +from typing import Any, Callable, Dict + +handlers: Dict[str, Callable[[bytes], Any]] = { + "decode": lambda data: "decoded" +} + +result = handlers["decode"](b"hello") + +reveal_type(result) # N: Revealed type is "builtins.str" +value: bytes = result # E: Incompatible types in assignment (expression has type "str", variable has type "bytes") + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py + +[case testUnionRuleReportsOnceAtDefinition] +# flags: --python-version 3.11 --config-file tmp/mypy.ini --show-error-codes + +def choose(value: int | str) -> int | str: # E: Union return types are not allowed [union-return] + return value + +def first_caller() -> None: + choose(1) + +def second_caller() -> None: + choose("hello") + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py + +[case testUnionRuleRunsWithoutAnyCalls] +# flags: --python-version 3.11 --config-file tmp/mypy.ini --show-error-codes + +def never_called(value: int | str) -> int | str: # E: Union return types are not allowed [union-return] + return value + + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py + +[case testUnionRuleCoversMethods] +# flags: --python-version 3.11 --config-file tmp/mypy.ini --show-error-codes + +class Parser: + def parse(self, value: int | str) -> int | str: # E: Union return types are not allowed [union-return] + return value + + +Parser().parse(1) +Parser().parse("hello") + + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py + +[case testUnionRuleCanBeIgnoredAtDefinition] +# flags: --python-version 3.11 --config-file tmp/mypy.ini --show-error-codes + +def permitted(value: int | str) -> int | str: # type: ignore[union-return] + return value + + +permitted(1) +permitted("hello") + + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py + +[case testAnyReturnRefinesEarlierAndLaterCallers] +# flags: --python-version 3.11 --config-file tmp/mypy.ini --show-error-codes + +from typing import Any + + +def before_definition() -> None: + reveal_type(bad(1)) # N: Revealed type is "builtins.str" + value: int = bad(1) # E: Incompatible types in assignment (expression has type "str", variable has type "int") [assignment] + + +def bad(x: int) -> Any: + return str(x) + + +def after_definition() -> None: + reveal_type(bad(1)) # N: Revealed type is "builtins.str" + value: int = bad(1) # E: Incompatible types in assignment (expression has type "str", variable has type "int") [assignment] + + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py + +[builtins fixtures/elaboration.pyi] + +[case testAnyReturnRefinementReachesFixedPoint] +# flags: --python-version 3.11 --config-file tmp/mypy.ini --show-error-codes + +from typing import Any + + +def use_result() -> None: + reveal_type(first()) # N: Revealed type is "Literal['done']?" + value: int = first() # E: Incompatible types in assignment (expression has type "str", variable has type "int") [assignment] + + +def first() -> Any: + return second() + + +def second() -> Any: + return "done" + + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py + +[case testAnyReturnRuleDoesNotChangeNonAnyAnnotation] +# flags: --python-version 3.11 --config-file tmp/mypy.ini --show-error-codes + +def returns_object() -> object: + return "hello" + + +reveal_type(returns_object()) # N: Revealed type is "builtins.object" + + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py + +[case testAnyReturnRuleLeavesBodyAnyAlone] +# flags: --python-version 3.11 --config-file tmp/mypy.ini --show-error-codes + +from typing import Any + + +def identity(value: Any) -> Any: + return value + + +reveal_type(identity(1)) # N: Revealed type is "Any" + + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/function_definition_rules.py diff --git a/test-data/unit/fixtures/elaboration.pyi b/test-data/unit/fixtures/elaboration.pyi new file mode 100644 index 0000000000000..ee18fffe6fadb --- /dev/null +++ b/test-data/unit/fixtures/elaboration.pyi @@ -0,0 +1,27 @@ +# Builtins stub for function-body elaboration tests. + +from typing import Any, Generic, TypeVar + +T = TypeVar("T") +K = TypeVar("K") +V = TypeVar("V") + +class object: + def __init__(self) -> None: pass + +class type: + def __init__(self, x: object) -> None: pass + +class int: pass +class bool(int): pass +class float: pass +class bytes: pass + +class str: + def __init__(self, object: object = "") -> None: pass + + +class tuple(Generic[T]): pass +class list(Generic[T]): pass +class dict(Generic[K, V]): pass +class function: pass diff --git a/test-data/unit/plugins/function_definition_rules.py b/test-data/unit/plugins/function_definition_rules.py new file mode 100644 index 0000000000000..184b03067c649 --- /dev/null +++ b/test-data/unit/plugins/function_definition_rules.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from collections.abc import Callable + +from mypy.errorcodes import ErrorCode +from mypy.plugin import ( + FunctionBodyContext, + FunctionBodyResult, + FunctionDefContext, + FunctionDefHookResult, + Plugin, +) +from mypy.types import AnyType, TypeOfAny, UnionType, get_proper_type + + +UNION_RETURN = ErrorCode( + code="union-return", + description="Disallow union return annotations", + category="General", +) + + +def refine_explicit_any_return(ctx: FunctionBodyContext) -> FunctionBodyResult | None: + inferred = get_proper_type(ctx.inferred_return_type) + if isinstance(inferred, AnyType): + return None + return FunctionBodyResult(refined_return=inferred) + + +def inspect_function_definition(ctx: FunctionDefContext) -> FunctionDefHookResult | None: + declared_return = get_proper_type(ctx.declared_signature.ret_type) + if isinstance(declared_return, UnionType): + ctx.api.fail( + "Union return types are not allowed", + ctx.definition, + code=UNION_RETURN, + ) + if ( + isinstance(declared_return, AnyType) + and declared_return.type_of_any == TypeOfAny.explicit + ): + return FunctionDefHookResult(after_body=refine_explicit_any_return) + return None + + +class FunctionDefinitionRulesPlugin(Plugin): + def get_function_def_hook( + self, fullname: str + ) -> Callable[[FunctionDefContext], FunctionDefHookResult | None] | None: + return inspect_function_definition + + +def plugin(version: str) -> type[FunctionDefinitionRulesPlugin]: + return FunctionDefinitionRulesPlugin