Implementation for validation checks on functions that can only be called on dataclasses

mypy/plugins/dataclasses.py

@@ -4,6 +4,7 @@
 
 from typing import Optional
 from typing_extensions import Final
+import mypy
 
 from mypy.expandtype import expand_type
 from mypy.nodes import (
@@ -630,6 +631,23 @@ def dataclass_class_maker_callback(ctx: ClassDefContext) -> bool:
     transformer = DataclassTransformer(ctx)
     return transformer.transform()
 
+def dataclass_check_type_callback(ctx: mypy.plugin.FunctionContext) -> Type:
+    """Will raise an error if the function is not called on a dataclass
+    instance. Examples of these functions are replace() and field() functions. 
+
+    This callback is called in default.py
+    """
+    if ctx != None:
+        if len(ctx.arg_types) != 0:
+            if len(ctx.arg_types[0]) != 0:
+                #this means there's no dataclass_tag in the metadata, so the argument is not a dataclass
+                if 'dataclass' not in (ctx.arg_types[0][0].type.metadata):
+                    ctx.api.msg.fail(
+                        "Calling replace on a non-dataclass.",
+                        ctx.context
+                    )
+    return ctx.default_return_type
+
 
 def _collect_field_args(
     expr: Expression, ctx: ClassDefContext

mypy/plugins/default.py

@@ -37,12 +37,15 @@ class DefaultPlugin(Plugin):
     """Type checker plugin that is enabled by default."""
 
     def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
-        from mypy.plugins import ctypes, singledispatch
+        from mypy.plugins import ctypes, singledispatch, dataclasses
 
         if fullname == "ctypes.Array":
             return ctypes.array_constructor_callback
         elif fullname == "functools.singledispatch":
             return singledispatch.create_singledispatch_function_callback
+        #these functions should only be called on dataclass instances
+        elif fullname == "dataclasses.replace" or fullname == "dataclasses.field":
+            return dataclasses.dataclass_check_type_callback
         return None
 
     def get_method_signature_hook(

test-data/unit/check-dataclasses.test

@@ -1,3 +1,48 @@
+[case testNoCrashOnReplace]
+def replace(__obj: _T, **changes: Any) -> _T: ...
+# flags: --python-version 3.7
+from dataclasses import dataclass
+from dataclasses import replace
+@dataclass
+class D:
+    pass
+class C:
+    pass
+
+replace(C())  # E: replace() should be called on dataclass instances
+[builtins fixtures/dataclasses.pyi]
+
+
+[case testNoCrashOnAstuple]
+def replace(__obj: _T, **changes: Any) -> _T: ...
+# flags: --python-version 3.7
+from dataclasses import dataclass
+from dataclasses import astuple
+
+@dataclass
+class D:
+    pass
+class C:
+    pass
+
+astuple(C())  # E: astuple() should be called on dataclass instances
+[builtins fixtures/dataclasses.pyi]
+
+[case testNoCrashOnFields]
+def replace(__obj: _T, **changes: Any) -> _T: ...
+# flags: --python-version 3.7
+from dataclasses import dataclass
+from dataclasses import fields
+
+@dataclass
+class D:
+    pass
+class C:
+    pass
+
+fields(C())   # E: astuple() should be called on dataclass instances
+[builtins fixtures/dataclasses.pyi]
+
 [case testDataclassesBasic]
 # flags: --python-version 3.7
 from dataclasses import dataclass

test-data/unit/lib-stub/dataclasses.pyi

@@ -30,5 +30,7 @@ def field(*,
     init: bool = ..., repr: bool = ..., hash: Optional[bool] = ..., compare: bool = ...,
     metadata: Optional[Mapping[str, Any]] = ..., kw_only: bool = ...,) -> Any: ...
 
+@overload 
+def replace(obj: _T, **changes: Any) -> _T: ...
 
 class Field(Generic[_T]): pass