Add signature for dataclasses.replace

mypy/plugins/dataclasses.py

@@ -7,6 +7,7 @@
 
 from mypy import errorcodes, message_registry
 from mypy.expandtype import expand_type
+from mypy.messages import format_type_bare
 from mypy.nodes import (
     ARG_NAMED,
     ARG_NAMED_OPT,
@@ -23,6 +24,7 @@
     Context,
     DataclassTransformSpec,
     Expression,
+    FuncDef,
     IfStmt,
     JsonDict,
     NameExpr,
@@ -37,7 +39,7 @@
     TypeVarExpr,
     Var,
 )
-from mypy.plugin import ClassDefContext, SemanticAnalyzerPluginInterface
+from mypy.plugin import ClassDefContext, FunctionSigContext, SemanticAnalyzerPluginInterface
 from mypy.plugins.common import (
     _get_decorator_bool_argument,
     add_attribute_to_class,
@@ -74,6 +76,7 @@
     frozen_default=False,
     field_specifiers=("dataclasses.Field", "dataclasses.field"),
 )
+_INTERNAL_REPLACE_SYM_NAME = "__mypy_replace"
 
 
 class DataclassAttribute:
@@ -326,6 +329,7 @@ def transform(self) -> bool:
             add_attribute_to_class(self._api, self._cls, "__match_args__", match_args_type)
 
         self._add_dataclass_fields_magic_attribute()
+        self._add_internal_replace_method(attributes)
 
         info.metadata["dataclass"] = {
             "attributes": [attr.serialize() for attr in attributes],
@@ -334,6 +338,35 @@ 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] = [Instance(self._cls.info, [])]
+        arg_kinds = [ARG_POS]
+        arg_names: list[str | None] = [None]
+        for attr in attributes:
+            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=Instance(self._cls.info, []),
+            fallback=self._api.named_type("builtins.function"),
+            name=f"replace of {self._cls.info.name}",
+        )
+
+        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:
@@ -787,3 +820,43 @@ def _is_dataclasses_decorator(node: Node) -> bool:
     if isinstance(node, RefExpr):
         return node.fullname in dataclass_makers
     return False
+
+
+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]
+
+    # <hack>
+    from mypy.checker import TypeChecker
+
+    assert isinstance(ctx.api, TypeChecker)
+    obj_type = ctx.api.expr_checker.accept(obj_arg)
+    # </hack>
+
+    obj_type = get_proper_type(obj_type)
+    if not isinstance(obj_type, Instance):
+        return ctx.default_signature
+
+    replace_func = obj_type.type.get_method(_INTERNAL_REPLACE_SYM_NAME)
+    if replace_func is None:
+        obj_type_str = format_type_bare(obj_type)
+        ctx.api.fail(
+            f'Argument 1 to "replace" has incompatible type "{obj_type_str}"; expected a dataclass',
+            ctx.context,
+        )
+        return ctx.default_signature
+
+    signature = get_proper_type(replace_func.type)
+    assert isinstance(signature, CallableType)
+    return signature

mypy/plugins/default.py

@@ -50,10 +50,12 @@ 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 == "dataclasses.replace":
+            return dataclasses.replace_function_sig_callback
         return None
 
     def get_method_signature_hook(

test-data/unit/check-dataclasses.test

@@ -2002,7 +2002,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
@@ -2036,4 +2035,38 @@ 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"
+
 [builtins fixtures/dataclasses.pyi]
+
+[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

test-data/unit/deps.test

@@ -1388,6 +1388,7 @@ class B(A):
 <m.A.(abstract)> -> <m.B.__init__>, m
 <m.A.__dataclass_fields__> -> <m.B.__dataclass_fields__>
 <m.A.__init__> -> <m.B.__init__>, m.B.__init__
+<m.A.__mypy_replace> -> <m.B.__mypy_replace>
 <m.A.__new__> -> <m.B.__new__>
 <m.A.x> -> <m.B.x>
 <m.A.y> -> <m.B.y>
@@ -1419,6 +1420,7 @@ class B(A):
 <m.A.__dataclass_fields__> -> <m.B.__dataclass_fields__>
 <m.A.__init__> -> <m.B.__init__>, m.B.__init__
 <m.A.__match_args__> -> <m.B.__match_args__>
+<m.A.__mypy_replace> -> <m.B.__mypy_replace>
 <m.A.__new__> -> <m.B.__new__>
 <m.A.x> -> <m.B.x>
 <m.A.y> -> <m.B.y>

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: ...

test-data/unit/pythoneval.test

@@ -1984,3 +1984,23 @@ def good9(foo1: Foo[Concatenate[int, P]], foo2: Foo[[int, str, bytes]], *args: P
 
 [out]
 _testStrictEqualitywithParamSpec.py:11: error: Non-overlapping equality check (left operand type: "Foo[[int]]", right operand type: "Bar[[int]]")
+
+[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"