diff --git a/mypy/plugins/common.py b/mypy/plugins/common.py index 0acf3e3a6369..67796ef15cf3 100644 --- a/mypy/plugins/common.py +++ b/mypy/plugins/common.py @@ -1,5 +1,6 @@ from __future__ import annotations +from mypy.argmap import map_actuals_to_formals from mypy.fixup import TypeFixer from mypy.nodes import ( ARG_POS, @@ -13,6 +14,7 @@ Expression, FuncDef, JsonDict, + NameExpr, Node, PassStmt, RefExpr, @@ -22,6 +24,7 @@ from mypy.plugin import CheckerPluginInterface, ClassDefContext, SemanticAnalyzerPluginInterface from mypy.semanal_shared import ( ALLOW_INCOMPATIBLE_OVERRIDE, + parse_bool, require_bool_literal_argument, set_callable_name, ) @@ -29,13 +32,19 @@ try_getting_str_literals as try_getting_str_literals, ) from mypy.types import ( + AnyType, CallableType, + Instance, + LiteralType, + NoneType, Overloaded, Type, + TypeOfAny, TypeType, TypeVarType, deserialize_type, get_proper_type, + is_optional, ) from mypy.typevars import fill_typevars from mypy.util import get_unique_redefinition_name @@ -87,6 +96,71 @@ def _get_argument(call: CallExpr, name: str) -> Expression | None: return None +def find_shallow_matching_overload_item(overload: Overloaded, call: CallExpr) -> CallableType: + """Perform limited lookup of a matching overload item. + + Full overload resolution is only supported during type checking, but plugins + sometimes need to resolve overloads. This can be used in some such use cases. + + Resolve overloads based on these things only: + + * Match using argument kinds and names + * If formal argument has type None, only accept the "None" expression in the callee + * If formal argument has type Literal[True] or Literal[False], only accept the + relevant bool literal + + Return the first matching overload item, or the last one if nothing matches. + """ + for item in overload.items[:-1]: + ok = True + mapped = map_actuals_to_formals( + call.arg_kinds, + call.arg_names, + item.arg_kinds, + item.arg_names, + lambda i: AnyType(TypeOfAny.special_form), + ) + + # Look for extra actuals + matched_actuals = set() + for actuals in mapped: + matched_actuals.update(actuals) + if any(i not in matched_actuals for i in range(len(call.args))): + ok = False + + for arg_type, kind, actuals in zip(item.arg_types, item.arg_kinds, mapped): + if kind.is_required() and not actuals: + # Missing required argument + ok = False + break + elif actuals: + args = [call.args[i] for i in actuals] + arg_type = get_proper_type(arg_type) + arg_none = any(isinstance(arg, NameExpr) and arg.name == "None" for arg in args) + if isinstance(arg_type, NoneType): + if not arg_none: + ok = False + break + elif ( + arg_none + and not is_optional(arg_type) + and not ( + isinstance(arg_type, Instance) + and arg_type.type.fullname == "builtins.object" + ) + and not isinstance(arg_type, AnyType) + ): + ok = False + break + elif isinstance(arg_type, LiteralType) and type(arg_type.value) is bool: + if not any(parse_bool(arg) == arg_type.value for arg in args): + ok = False + break + if ok: + return item + return overload.items[-1] + + def _get_callee_type(call: CallExpr) -> CallableType | None: """Return the type of the callee, regardless of its syntatic form.""" @@ -103,8 +177,7 @@ def _get_callee_type(call: CallExpr) -> CallableType | None: if isinstance(callee_node, (Var, SYMBOL_FUNCBASE_TYPES)) and callee_node.type: callee_node_type = get_proper_type(callee_node.type) if isinstance(callee_node_type, Overloaded): - # We take the last overload. - return callee_node_type.items[-1] + return find_shallow_matching_overload_item(callee_node_type, call) elif isinstance(callee_node_type, CallableType): return callee_node_type diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py index e84e1dbb9491..9fa32a108b74 100644 --- a/mypy/plugins/dataclasses.py +++ b/mypy/plugins/dataclasses.py @@ -40,6 +40,7 @@ ) from mypy.plugin import ClassDefContext, SemanticAnalyzerPluginInterface from mypy.plugins.common import ( + _get_callee_type, _get_decorator_bool_argument, add_attribute_to_class, add_method_to_class, @@ -48,7 +49,7 @@ from mypy.semanal_shared import find_dataclass_transform_spec, require_bool_literal_argument from mypy.server.trigger import make_wildcard_trigger from mypy.state import state -from mypy.typeops import map_type_from_supertype +from mypy.typeops import map_type_from_supertype, try_getting_literals_from_type from mypy.types import ( AnyType, CallableType, @@ -517,7 +518,7 @@ def collect_attributes(self) -> list[DataclassAttribute] | None: is_in_init_param = field_args.get("init") if is_in_init_param is None: - is_in_init = True + is_in_init = self._get_default_init_value_for_field_specifier(stmt.rvalue) else: is_in_init = bool(self._api.parse_bool(is_in_init_param)) @@ -760,6 +761,30 @@ def _get_bool_arg(self, name: str, default: bool) -> bool: return require_bool_literal_argument(self._api, expression, name, default) return default + def _get_default_init_value_for_field_specifier(self, call: Expression) -> bool: + """ + Find a default value for the `init` parameter of the specifier being called. If the + specifier's type signature includes an `init` parameter with a type of `Literal[True]` or + `Literal[False]`, return the appropriate boolean value from the literal. Otherwise, + fall back to the standard default of `True`. + """ + if not isinstance(call, CallExpr): + return True + + specifier_type = _get_callee_type(call) + if specifier_type is None: + return True + + parameter = specifier_type.argument_by_name("init") + if parameter is None: + return True + + literals = try_getting_literals_from_type(parameter.typ, bool, "builtins.bool") + if literals is None or len(literals) != 1: + return True + + return literals[0] + def _infer_dataclass_attr_init_type( self, sym: SymbolTableNode, name: str, context: Context ) -> Type | None: diff --git a/mypy/semanal.py b/mypy/semanal.py index 4ee18d5ff4d3..84573790fd68 100644 --- a/mypy/semanal.py +++ b/mypy/semanal.py @@ -216,6 +216,7 @@ calculate_tuple_fallback, find_dataclass_transform_spec, has_placeholder, + parse_bool, require_bool_literal_argument, set_callable_name as set_callable_name, ) @@ -6462,12 +6463,8 @@ def is_initial_mangled_global(self, name: str) -> bool: return name == unmangle(name) + "'" def parse_bool(self, expr: Expression) -> bool | None: - if isinstance(expr, NameExpr): - if expr.fullname == "builtins.True": - return True - if expr.fullname == "builtins.False": - return False - return None + # This wrapper is preserved for plugins. + return parse_bool(expr) def parse_str_literal(self, expr: Expression) -> str | None: """Attempt to find the string literal value of the given expression. Returns `None` if no diff --git a/mypy/semanal_shared.py b/mypy/semanal_shared.py index 03efbe6ca1b8..c86ed828b2b9 100644 --- a/mypy/semanal_shared.py +++ b/mypy/semanal_shared.py @@ -18,6 +18,7 @@ Decorator, Expression, FuncDef, + NameExpr, Node, OverloadedFuncDef, RefExpr, @@ -451,7 +452,7 @@ def require_bool_literal_argument( default: bool | None = None, ) -> bool | None: """Attempt to interpret an expression as a boolean literal, and fail analysis if we can't.""" - value = api.parse_bool(expression) + value = parse_bool(expression) if value is None: api.fail( f'"{name}" argument must be a True or False literal', expression, code=LITERAL_REQ @@ -459,3 +460,12 @@ def require_bool_literal_argument( return default return value + + +def parse_bool(expr: Expression) -> bool | None: + if isinstance(expr, NameExpr): + if expr.fullname == "builtins.True": + return True + if expr.fullname == "builtins.False": + return False + return None diff --git a/mypy/test/testtypes.py b/mypy/test/testtypes.py index ee0256e2057a..6fe65675554b 100644 --- a/mypy/test/testtypes.py +++ b/mypy/test/testtypes.py @@ -7,7 +7,21 @@ from mypy.indirection import TypeIndirectionVisitor from mypy.join import join_simple, join_types from mypy.meet import meet_types, narrow_declared_type -from mypy.nodes import ARG_OPT, ARG_POS, ARG_STAR, ARG_STAR2, CONTRAVARIANT, COVARIANT, INVARIANT +from mypy.nodes import ( + ARG_NAMED, + ARG_OPT, + ARG_POS, + ARG_STAR, + ARG_STAR2, + CONTRAVARIANT, + COVARIANT, + INVARIANT, + ArgKind, + CallExpr, + Expression, + NameExpr, +) +from mypy.plugins.common import find_shallow_matching_overload_item from mypy.state import state from mypy.subtypes import is_more_precise, is_proper_subtype, is_same_type, is_subtype from mypy.test.helpers import Suite, assert_equal, assert_type, skip @@ -1287,3 +1301,135 @@ def assert_union_result(self, t: ProperType, expected: list[Type]) -> None: t2 = remove_instance_last_known_values(t) assert type(t2) is UnionType assert t2.items == expected + + +class ShallowOverloadMatchingSuite(Suite): + def setUp(self) -> None: + self.fx = TypeFixture() + + def test_simple(self) -> None: + fx = self.fx + ov = self.make_overload([[("x", fx.anyt, ARG_NAMED)], [("y", fx.anyt, ARG_NAMED)]]) + # Match first only + self.assert_find_shallow_matching_overload_item(ov, make_call(("foo", "x")), 0) + # Match second only + self.assert_find_shallow_matching_overload_item(ov, make_call(("foo", "y")), 1) + # No match -- invalid keyword arg name + self.assert_find_shallow_matching_overload_item(ov, make_call(("foo", "z")), 1) + # No match -- missing arg + self.assert_find_shallow_matching_overload_item(ov, make_call(), 1) + # No match -- extra arg + self.assert_find_shallow_matching_overload_item( + ov, make_call(("foo", "x"), ("foo", "z")), 1 + ) + + def test_match_using_types(self) -> None: + fx = self.fx + ov = self.make_overload( + [ + [("x", fx.nonet, ARG_POS)], + [("x", fx.lit_false, ARG_POS)], + [("x", fx.lit_true, ARG_POS)], + [("x", fx.anyt, ARG_POS)], + ] + ) + self.assert_find_shallow_matching_overload_item(ov, make_call(("None", None)), 0) + self.assert_find_shallow_matching_overload_item(ov, make_call(("builtins.False", None)), 1) + self.assert_find_shallow_matching_overload_item(ov, make_call(("builtins.True", None)), 2) + self.assert_find_shallow_matching_overload_item(ov, make_call(("foo", None)), 3) + + def test_none_special_cases(self) -> None: + fx = self.fx + ov = self.make_overload( + [[("x", fx.callable(fx.nonet), ARG_POS)], [("x", fx.nonet, ARG_POS)]] + ) + self.assert_find_shallow_matching_overload_item(ov, make_call(("None", None)), 1) + self.assert_find_shallow_matching_overload_item(ov, make_call(("func", None)), 0) + ov = self.make_overload([[("x", fx.str_type, ARG_POS)], [("x", fx.nonet, ARG_POS)]]) + self.assert_find_shallow_matching_overload_item(ov, make_call(("None", None)), 1) + self.assert_find_shallow_matching_overload_item(ov, make_call(("func", None)), 0) + ov = self.make_overload( + [[("x", UnionType([fx.str_type, fx.a]), ARG_POS)], [("x", fx.nonet, ARG_POS)]] + ) + self.assert_find_shallow_matching_overload_item(ov, make_call(("None", None)), 1) + self.assert_find_shallow_matching_overload_item(ov, make_call(("func", None)), 0) + ov = self.make_overload([[("x", fx.o, ARG_POS)], [("x", fx.nonet, ARG_POS)]]) + self.assert_find_shallow_matching_overload_item(ov, make_call(("None", None)), 0) + self.assert_find_shallow_matching_overload_item(ov, make_call(("func", None)), 0) + ov = self.make_overload( + [[("x", UnionType([fx.str_type, fx.nonet]), ARG_POS)], [("x", fx.nonet, ARG_POS)]] + ) + self.assert_find_shallow_matching_overload_item(ov, make_call(("None", None)), 0) + self.assert_find_shallow_matching_overload_item(ov, make_call(("func", None)), 0) + ov = self.make_overload([[("x", fx.anyt, ARG_POS)], [("x", fx.nonet, ARG_POS)]]) + self.assert_find_shallow_matching_overload_item(ov, make_call(("None", None)), 0) + self.assert_find_shallow_matching_overload_item(ov, make_call(("func", None)), 0) + + def test_optional_arg(self) -> None: + fx = self.fx + ov = self.make_overload( + [[("x", fx.anyt, ARG_NAMED)], [("y", fx.anyt, ARG_OPT)], [("z", fx.anyt, ARG_NAMED)]] + ) + self.assert_find_shallow_matching_overload_item(ov, make_call(), 1) + self.assert_find_shallow_matching_overload_item(ov, make_call(("foo", "x")), 0) + self.assert_find_shallow_matching_overload_item(ov, make_call(("foo", "y")), 1) + self.assert_find_shallow_matching_overload_item(ov, make_call(("foo", "z")), 2) + + def test_two_args(self) -> None: + fx = self.fx + ov = self.make_overload( + [ + [("x", fx.nonet, ARG_OPT), ("y", fx.anyt, ARG_OPT)], + [("x", fx.anyt, ARG_OPT), ("y", fx.anyt, ARG_OPT)], + ] + ) + self.assert_find_shallow_matching_overload_item(ov, make_call(), 0) + self.assert_find_shallow_matching_overload_item(ov, make_call(("None", "x")), 0) + self.assert_find_shallow_matching_overload_item(ov, make_call(("foo", "x")), 1) + self.assert_find_shallow_matching_overload_item( + ov, make_call(("foo", "y"), ("None", "x")), 0 + ) + self.assert_find_shallow_matching_overload_item( + ov, make_call(("foo", "y"), ("bar", "x")), 1 + ) + + def assert_find_shallow_matching_overload_item( + self, ov: Overloaded, call: CallExpr, expected_index: int + ) -> None: + c = find_shallow_matching_overload_item(ov, call) + assert c in ov.items + assert ov.items.index(c) == expected_index + + def make_overload(self, items: list[list[tuple[str, Type, ArgKind]]]) -> Overloaded: + result = [] + for item in items: + arg_types = [] + arg_names = [] + arg_kinds = [] + for name, typ, kind in item: + arg_names.append(name) + arg_types.append(typ) + arg_kinds.append(kind) + result.append( + CallableType( + arg_types, arg_kinds, arg_names, ret_type=NoneType(), fallback=self.fx.o + ) + ) + return Overloaded(result) + + +def make_call(*items: tuple[str, str | None]) -> CallExpr: + args: list[Expression] = [] + arg_names = [] + arg_kinds = [] + for arg, name in items: + shortname = arg.split(".")[-1] + n = NameExpr(shortname) + n.fullname = arg + args.append(n) + arg_names.append(name) + if name: + arg_kinds.append(ARG_NAMED) + else: + arg_kinds.append(ARG_POS) + return CallExpr(NameExpr("f"), args, arg_kinds, arg_names) diff --git a/mypy/test/typefixture.py b/mypy/test/typefixture.py index d12e7abab0e2..1013b87c213f 100644 --- a/mypy/test/typefixture.py +++ b/mypy/test/typefixture.py @@ -136,6 +136,7 @@ def make_type_var( self.type_type = Instance(self.type_typei, []) # type self.function = Instance(self.functioni, []) # function TODO self.str_type = Instance(self.str_type_info, []) + self.bool_type = Instance(self.bool_type_info, []) self.a = Instance(self.ai, []) # A self.b = Instance(self.bi, []) # B self.c = Instance(self.ci, []) # C @@ -197,6 +198,9 @@ def make_type_var( self.lit_str2_inst = Instance(self.str_type_info, [], last_known_value=self.lit_str2) self.lit_str3_inst = Instance(self.str_type_info, [], last_known_value=self.lit_str3) + self.lit_false = LiteralType(False, self.bool_type) + self.lit_true = LiteralType(True, self.bool_type) + self.type_a = TypeType.make_normalized(self.a) self.type_b = TypeType.make_normalized(self.b) self.type_c = TypeType.make_normalized(self.c) diff --git a/test-data/unit/check-dataclass-transform.test b/test-data/unit/check-dataclass-transform.test index 0b696d36b9bd..be6b46d70846 100644 --- a/test-data/unit/check-dataclass-transform.test +++ b/test-data/unit/check-dataclass-transform.test @@ -328,6 +328,38 @@ Foo(a=1, b='bye') [typing fixtures/typing-full.pyi] [builtins fixtures/dataclasses.pyi] +[case testDataclassTransformFieldSpecifierImplicitInit] +# flags: --python-version 3.11 +from typing import dataclass_transform, Literal, overload + +def init(*, init: Literal[True] = True): ... +def no_init(*, init: Literal[False] = False): ... + +@overload +def field_overload(*, custom: None, init: Literal[True] = True): ... +@overload +def field_overload(*, custom: str, init: Literal[False] = False): ... +def field_overload(*, custom, init): ... + +@dataclass_transform(field_specifiers=(init, no_init, field_overload)) +def my_dataclass(cls): return cls + +@my_dataclass +class Foo: + a: int = init() + b: int = field_overload(custom=None) + + bad1: int = no_init() + bad2: int = field_overload(custom="bad2") + +reveal_type(Foo) # N: Revealed type is "def (a: builtins.int, b: builtins.int) -> __main__.Foo" +Foo(a=1, b=2) +Foo(a=1, b=2, bad1=0) # E: Unexpected keyword argument "bad1" for "Foo" +Foo(a=1, b=2, bad2=0) # E: Unexpected keyword argument "bad2" for "Foo" + +[typing fixtures/typing-full.pyi] +[builtins fixtures/dataclasses.pyi] + [case testDataclassTransformOverloadsDecoratorOnOverload] # flags: --python-version 3.11 from typing import dataclass_transform, overload, Any, Callable, Type, Literal