diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py index 2fd903f2f8b9..913b1ef23312 100644 --- a/mypy/plugins/dataclasses.py +++ b/mypy/plugins/dataclasses.py @@ -7,6 +7,8 @@ from mypy import errorcodes, message_registry from mypy.expandtype import expand_type, expand_type_by_instance +from mypy.meet import meet_types +from mypy.messages import format_type_bare from mypy.nodes import ( ARG_NAMED, ARG_NAMED_OPT, @@ -38,7 +40,7 @@ TypeVarExpr, Var, ) -from mypy.plugin import ClassDefContext, SemanticAnalyzerPluginInterface +from mypy.plugin import ClassDefContext, FunctionSigContext, SemanticAnalyzerPluginInterface from mypy.plugins.common import ( _get_callee_type, _get_decorator_bool_argument, @@ -56,10 +58,13 @@ Instance, LiteralType, NoneType, + ProperType, TupleType, Type, TypeOfAny, TypeVarType, + UninhabitedType, + UnionType, get_proper_type, ) from mypy.typevars import fill_typevars @@ -76,6 +81,7 @@ frozen_default=False, field_specifiers=("dataclasses.Field", "dataclasses.field"), ) +_INTERNAL_REPLACE_SYM_NAME = "__mypy-replace" class DataclassAttribute: @@ -344,6 +350,9 @@ def transform(self) -> bool: self._add_dataclass_fields_magic_attribute() + if self._spec is _TRANSFORM_SPEC_FOR_DATACLASSES: + self._add_internal_replace_method(attributes) + info.metadata["dataclass"] = { "attributes": [attr.serialize() for attr in attributes], "frozen": decorator_arguments["frozen"], @@ -351,6 +360,37 @@ def transform(self) -> bool: return True + def _add_internal_replace_method(self, attributes: list[DataclassAttribute]) -> None: + """ + Stashes the signature of 'dataclasses.replace(...)' for this specific dataclass + to be used later whenever 'dataclasses.replace' is called for this dataclass. + """ + arg_types: list[Type] = [] + arg_kinds = [] + arg_names: list[str | None] = [] + + info = self._cls.info + for attr in attributes: + attr_type = attr.expand_type(info) + assert attr_type is not None + arg_types.append(attr_type) + arg_kinds.append( + ARG_NAMED if attr.is_init_var and not attr.has_default else ARG_NAMED_OPT + ) + arg_names.append(attr.name) + + signature = CallableType( + arg_types=arg_types, + arg_kinds=arg_kinds, + arg_names=arg_names, + ret_type=NoneType(), + fallback=self._api.named_type("builtins.function"), + ) + + self._cls.info.names[_INTERNAL_REPLACE_SYM_NAME] = SymbolTableNode( + kind=MDEF, node=FuncDef(typ=signature), plugin_generated=True + ) + def add_slots( self, info: TypeInfo, attributes: list[DataclassAttribute], *, correct_version: bool ) -> None: @@ -893,3 +933,124 @@ def _has_direct_dataclass_transform_metaclass(info: TypeInfo) -> bool: info.declared_metaclass is not None and info.declared_metaclass.type.dataclass_transform_spec is not None ) + + +def _fail_not_dataclass(ctx: FunctionSigContext, t: Type, parent_t: Type) -> None: + t_name = format_type_bare(t, ctx.api.options) + if parent_t is t: + msg = ( + f'Argument 1 to "replace" has a variable type "{t_name}" not bound to a dataclass' + if isinstance(t, TypeVarType) + else f'Argument 1 to "replace" has incompatible type "{t_name}"; expected a dataclass' + ) + else: + pt_name = format_type_bare(parent_t, ctx.api.options) + msg = ( + f'Argument 1 to "replace" has type "{pt_name}" whose item "{t_name}" is not bound to a dataclass' + if isinstance(t, TypeVarType) + else f'Argument 1 to "replace" has incompatible type "{pt_name}" whose item "{t_name}" is not a dataclass' + ) + + ctx.api.fail(msg, ctx.context) + + +def _get_expanded_dataclasses_fields( + ctx: FunctionSigContext, typ: ProperType, display_typ: ProperType, parent_typ: ProperType +) -> list[CallableType] | None: + """ + For a given type, determine what dataclasses it can be: for each class, return the field types. + For generic classes, the field types are expanded. + If the type contains Any or a non-dataclass, returns None; in the latter case, also reports an error. + """ + if isinstance(typ, AnyType): + return None + elif isinstance(typ, UnionType): + ret: list[CallableType] | None = [] + for item in typ.relevant_items(): + item = get_proper_type(item) + item_types = _get_expanded_dataclasses_fields(ctx, item, item, parent_typ) + if ret is not None and item_types is not None: + ret += item_types + else: + ret = None # but keep iterating to emit all errors + return ret + elif isinstance(typ, TypeVarType): + return _get_expanded_dataclasses_fields( + ctx, get_proper_type(typ.upper_bound), display_typ, parent_typ + ) + elif isinstance(typ, Instance): + replace_sym = typ.type.get_method(_INTERNAL_REPLACE_SYM_NAME) + if replace_sym is None: + _fail_not_dataclass(ctx, display_typ, parent_typ) + return None + replace_sig = replace_sym.type + assert isinstance(replace_sig, ProperType) + assert isinstance(replace_sig, CallableType) + return [expand_type_by_instance(replace_sig, typ)] + else: + _fail_not_dataclass(ctx, display_typ, parent_typ) + return None + + +# TODO: we can potentially get the function signature hook to allow returning a union +# and leave this to the regular machinery of resolving a union of callables +# (https://github.com/python/mypy/issues/15457) +def _meet_replace_sigs(sigs: list[CallableType]) -> CallableType: + """ + Produces the lowest bound of the 'replace' signatures of multiple dataclasses. + """ + args = { + name: (typ, kind) + for name, typ, kind in zip(sigs[0].arg_names, sigs[0].arg_types, sigs[0].arg_kinds) + } + + for sig in sigs[1:]: + sig_args = { + name: (typ, kind) + for name, typ, kind in zip(sig.arg_names, sig.arg_types, sig.arg_kinds) + } + for name in (*args.keys(), *sig_args.keys()): + sig_typ, sig_kind = args.get(name, (UninhabitedType(), ARG_NAMED_OPT)) + sig2_typ, sig2_kind = sig_args.get(name, (UninhabitedType(), ARG_NAMED_OPT)) + args[name] = ( + meet_types(sig_typ, sig2_typ), + ARG_NAMED_OPT if sig_kind == sig2_kind == ARG_NAMED_OPT else ARG_NAMED, + ) + + return sigs[0].copy_modified( + arg_names=list(args.keys()), + arg_types=[typ for typ, _ in args.values()], + arg_kinds=[kind for _, kind in args.values()], + ) + + +def replace_function_sig_callback(ctx: FunctionSigContext) -> CallableType: + """ + Returns a signature for the 'dataclasses.replace' function that's dependent on the type + of the first positional argument. + """ + if len(ctx.args) != 2: + # Ideally the name and context should be callee's, but we don't have it in FunctionSigContext. + ctx.api.fail(f'"{ctx.default_signature.name}" has unexpected type annotation', ctx.context) + return ctx.default_signature + + if len(ctx.args[0]) != 1: + return ctx.default_signature # leave it to the type checker to complain + + obj_arg = ctx.args[0][0] + obj_type = get_proper_type(ctx.api.get_expression_type(obj_arg)) + inst_type_str = format_type_bare(obj_type, ctx.api.options) + + replace_sigs = _get_expanded_dataclasses_fields(ctx, obj_type, obj_type, obj_type) + if replace_sigs is None: + return ctx.default_signature + replace_sig = _meet_replace_sigs(replace_sigs) + + return replace_sig.copy_modified( + arg_names=[None, *replace_sig.arg_names], + arg_kinds=[ARG_POS, *replace_sig.arg_kinds], + arg_types=[obj_type, *replace_sig.arg_types], + ret_type=obj_type, + fallback=ctx.default_signature.fallback, + name=f"{ctx.default_signature.name} of {inst_type_str}", + ) diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index 500eef76a9d9..b83c0192a14b 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -51,12 +51,14 @@ def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] def get_function_signature_hook( self, fullname: str ) -> Callable[[FunctionSigContext], FunctionLike] | None: - from mypy.plugins import attrs + from mypy.plugins import attrs, dataclasses if fullname in ("attr.evolve", "attrs.evolve", "attr.assoc", "attrs.assoc"): return attrs.evolve_function_sig_callback elif fullname in ("attr.fields", "attrs.fields"): return attrs.fields_function_sig_callback + elif fullname == "dataclasses.replace": + return dataclasses.replace_function_sig_callback return None def get_method_signature_hook( diff --git a/test-data/unit/check-dataclass-transform.test b/test-data/unit/check-dataclass-transform.test index be6b46d70846..53a6eea95bfa 100644 --- a/test-data/unit/check-dataclass-transform.test +++ b/test-data/unit/check-dataclass-transform.test @@ -840,6 +840,24 @@ reveal_type(bar.base) # N: Revealed type is "builtins.int" [typing fixtures/typing-full.pyi] [builtins fixtures/dataclasses.pyi] +[case testDataclassTransformReplace] +from dataclasses import replace +from typing import dataclass_transform, Type + +@dataclass_transform() +def my_dataclass(cls: Type) -> Type: + return cls + +@my_dataclass +class Person: + name: str + +p = Person('John') +y = replace(p, name='Bob') # E: Argument 1 to "replace" has incompatible type "Person"; expected a dataclass + +[typing fixtures/typing-full.pyi] +[builtins fixtures/dataclasses.pyi] + [case testDataclassTransformSimpleDescriptor] # flags: --python-version 3.11 @@ -1051,5 +1069,6 @@ class Desc2: class C: x: Desc # E: Unsupported signature for "__set__" in "Desc" y: Desc2 # E: Unsupported "__set__" in "Desc2" + [typing fixtures/typing-full.pyi] [builtins fixtures/dataclasses.pyi] diff --git a/test-data/unit/check-dataclasses.test b/test-data/unit/check-dataclasses.test index 914e1c2e0602..1f6c8d143243 100644 --- a/test-data/unit/check-dataclasses.test +++ b/test-data/unit/check-dataclasses.test @@ -2006,7 +2006,6 @@ e: Element[Bar] reveal_type(e.elements) # N: Revealed type is "typing.Sequence[__main__.Element[__main__.Bar]]" [builtins fixtures/dataclasses.pyi] - [case testIfConditionsInDefinition] # flags: --python-version 3.11 --always-true TRUTH from dataclasses import dataclass @@ -2040,8 +2039,164 @@ Foo( present_4=4, present_5=5, ) + +[builtins fixtures/dataclasses.pyi] + +[case testReplace] +from dataclasses import dataclass, replace, InitVar +from typing import ClassVar + +@dataclass +class A: + x: int + q: InitVar[int] + q2: InitVar[int] = 0 + c: ClassVar[int] + + +a = A(x=42, q=7) +a2 = replace(a) # E: Missing named argument "q" for "replace" of "A" +a2 = replace(a, q=42) +a2 = replace(a, x=42, q=42) +a2 = replace(a, x=42, q=42, c=7) # E: Unexpected keyword argument "c" for "replace" of "A" +a2 = replace(a, x='42', q=42) # E: Argument "x" to "replace" of "A" has incompatible type "str"; expected "int" +a2 = replace(a, q='42') # E: Argument "q" to "replace" of "A" has incompatible type "str"; expected "int" +reveal_type(a2) # N: Revealed type is "__main__.A" + +[case testReplaceUnion] +# flags: --strict-optional +from typing import Generic, Union, TypeVar +from dataclasses import dataclass, replace, InitVar + +T = TypeVar('T') + +@dataclass +class A(Generic[T]): + x: T # exercises meet(T=int, int) = int + y: bool # exercises meet(bool, int) = bool + z: str # exercises meet(str, bytes) = + w: dict # exercises meet(dict, ) = + init_var: InitVar[int] # exercises (non-optional, optional) = non-optional + +@dataclass +class B: + x: int + y: int + z: bytes + init_var: int + + +a_or_b: Union[A[int], B] +_ = replace(a_or_b, x=42, y=True, init_var=42) +_ = replace(a_or_b, x=42, y=True) # E: Missing named argument "init_var" for "replace" of "Union[A[int], B]" +_ = replace(a_or_b, x=42, y=True, z='42', init_var=42) # E: Argument "z" to "replace" of "Union[A[int], B]" has incompatible type "str"; expected +_ = replace(a_or_b, x=42, y=True, w={}, init_var=42) # E: Argument "w" to "replace" of "Union[A[int], B]" has incompatible type "Dict[, ]"; expected +_ = replace(a_or_b, y=42, init_var=42) # E: Argument "y" to "replace" of "Union[A[int], B]" has incompatible type "int"; expected "bool" + [builtins fixtures/dataclasses.pyi] +[case testReplaceUnionOfTypeVar] +# flags: --strict-optional +from typing import Generic, Union, TypeVar +from dataclasses import dataclass, replace + +@dataclass +class A: + x: int + y: int + z: str + w: dict + +class B: + pass + +TA = TypeVar('TA', bound=A) +TB = TypeVar('TB', bound=B) + +def f(b_or_t: Union[TA, TB, int]) -> None: + a2 = replace(b_or_t) # E: Argument 1 to "replace" has type "Union[TA, TB, int]" whose item "TB" is not bound to a dataclass # E: Argument 1 to "replace" has incompatible type "Union[TA, TB, int]" whose item "int" is not a dataclass + +[case testReplaceTypeVarBoundNotDataclass] +from dataclasses import dataclass, replace +from typing import Union, TypeVar + +TInt = TypeVar('TInt', bound=int) +TAny = TypeVar('TAny') +TNone = TypeVar('TNone', bound=None) +TUnion = TypeVar('TUnion', bound=Union[str, int]) + +def f1(t: TInt) -> None: + _ = replace(t, x=42) # E: Argument 1 to "replace" has a variable type "TInt" not bound to a dataclass + +def f2(t: TAny) -> TAny: + return replace(t, x='spam') # E: Argument 1 to "replace" has a variable type "TAny" not bound to a dataclass + +def f3(t: TNone) -> TNone: + return replace(t, x='spam') # E: Argument 1 to "replace" has a variable type "TNone" not bound to a dataclass + +def f4(t: TUnion) -> TUnion: + return replace(t, x='spam') # E: Argument 1 to "replace" has incompatible type "TUnion" whose item "str" is not a dataclass # E: Argument 1 to "replace" has incompatible type "TUnion" whose item "int" is not a dataclass + +[case testReplaceTypeVarBound] +from dataclasses import dataclass, replace +from typing import TypeVar + +@dataclass +class A: + x: int + +@dataclass +class B(A): + pass + +TA = TypeVar('TA', bound=A) + +def f(t: TA) -> TA: + t2 = replace(t, x=42) + reveal_type(t2) # N: Revealed type is "TA`-1" + _ = replace(t, x='42') # E: Argument "x" to "replace" of "TA" has incompatible type "str"; expected "int" + return t2 + +f(A(x=42)) +f(B(x=42)) + +[case testReplaceAny] +from dataclasses import replace +from typing import Any + +a: Any +a2 = replace(a) +reveal_type(a2) # N: Revealed type is "Any" + +[case testReplaceNotDataclass] +from dataclasses import replace + +replace(5) # E: Argument 1 to "replace" has incompatible type "int"; expected a dataclass + +class C: + pass + +replace(C()) # E: Argument 1 to "replace" has incompatible type "C"; expected a dataclass + +replace(None) # E: Argument 1 to "replace" has incompatible type "None"; expected a dataclass + +[case testReplaceGeneric] +from dataclasses import dataclass, replace, InitVar +from typing import ClassVar, Generic, TypeVar + +T = TypeVar('T') + +@dataclass +class A(Generic[T]): + x: T + +a = A(x=42) +reveal_type(a) # N: Revealed type is "__main__.A[builtins.int]" +a2 = replace(a, x=42) +reveal_type(a2) # N: Revealed type is "__main__.A[builtins.int]" +a2 = replace(a, x='42') # E: Argument "x" to "replace" of "A[int]" has incompatible type "str"; expected "int" +reveal_type(a2) # N: Revealed type is "__main__.A[builtins.int]" + [case testProtocolNoCrash] from typing import Protocol, Union, ClassVar from dataclasses import dataclass, field diff --git a/test-data/unit/deps.test b/test-data/unit/deps.test index 28d51f1a4c30..fe5107b1529d 100644 --- a/test-data/unit/deps.test +++ b/test-data/unit/deps.test @@ -1388,6 +1388,7 @@ class B(A): -> , m -> -> , m.B.__init__ + -> -> -> -> @@ -1419,6 +1420,7 @@ class B(A): -> -> , m.B.__init__ -> + -> -> -> -> diff --git a/test-data/unit/fine-grained-dataclass.test b/test-data/unit/fine-grained-dataclass.test new file mode 100644 index 000000000000..036d858ddf69 --- /dev/null +++ b/test-data/unit/fine-grained-dataclass.test @@ -0,0 +1,25 @@ +[case testReplace] +[file model.py] +from dataclasses import dataclass + +@dataclass +class Model: + x: int = 0 +[file replace.py] +from dataclasses import replace +from model import Model + +m = Model() +replace(m, x=42) + +[file model.py.2] +from dataclasses import dataclass + +@dataclass +class Model: + x: str = 'hello' + +[builtins fixtures/dataclasses.pyi] +[out] +== +replace.py:5: error: Argument "x" to "replace" of "Model" has incompatible type "int"; expected "str" diff --git a/test-data/unit/lib-stub/dataclasses.pyi b/test-data/unit/lib-stub/dataclasses.pyi index bd33b459266c..b2b48c2ae486 100644 --- a/test-data/unit/lib-stub/dataclasses.pyi +++ b/test-data/unit/lib-stub/dataclasses.pyi @@ -32,3 +32,5 @@ def field(*, class Field(Generic[_T]): pass + +def replace(__obj: _T, **changes: Any) -> _T: ... diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test index 93ff2da9f7b9..8c82b6843a3b 100644 --- a/test-data/unit/pythoneval.test +++ b/test-data/unit/pythoneval.test @@ -2024,3 +2024,23 @@ def foo(callback: Callable[[], Any]) -> None: def call(callback: Callable[[Unpack[Ts]], Any], *args: Unpack[Ts]) -> Any: ... + +[case testDataclassReplace] +from dataclasses import dataclass, replace + +@dataclass +class A: + x: int + + +a = A(x=42) +a2 = replace(a, x=42) +reveal_type(a2) +a2 = replace() +a2 = replace(a, x='spam') +a2 = replace(a, x=42, q=42) +[out] +_testDataclassReplace.py:10: note: Revealed type is "_testDataclassReplace.A" +_testDataclassReplace.py:11: error: Too few arguments for "replace" +_testDataclassReplace.py:12: error: Argument "x" to "replace" of "A" has incompatible type "str"; expected "int" +_testDataclassReplace.py:13: error: Unexpected keyword argument "q" for "replace" of "A"