Generic dataclass support

ikonst · ikonst · commit c0058950753e · 2023-03-30T18:04:59.000-04:00
diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py
@@ -6,7 +6,7 @@
 from typing_extensions import Final
 
 from mypy import errorcodes, message_registry
-from mypy.expandtype import expand_type
+from mypy.expandtype import expand_type, expand_type_by_instance
 from mypy.messages import format_type_bare
 from mypy.nodes import (
     ARG_NAMED,
@@ -350,12 +350,15 @@ def _add_internal_replace_method(self, attributes: list[DataclassAttribute]) ->
         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]
+        arg_types: list[Type] = []
+        arg_kinds = []
+        arg_names: list[str | None] = []
+
+        info = self._cls.info
         for attr in attributes:
-            assert attr.type is not None
-            arg_types.append(attr.type)
+            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
             )
@@ -365,7 +368,7 @@ def _add_internal_replace_method(self, attributes: list[DataclassAttribute]) ->
             arg_types=arg_types,
             arg_kinds=arg_kinds,
             arg_names=arg_names,
-            ret_type=Instance(self._cls.info, []),
+            ret_type=NoneType(),
             fallback=self._api.named_type("builtins.function"),
             name=f"replace of {self._cls.info.name}",
         )
@@ -883,4 +886,12 @@ def replace_function_sig_callback(ctx: FunctionSigContext) -> CallableType:
 
     signature = get_proper_type(replace_func.type)
     assert isinstance(signature, CallableType)
+    signature = expand_type_by_instance(signature, obj_type)
+    # re-add the instance type
+    signature = signature.copy_modified(
+        arg_types=[obj_type, *signature.arg_types],
+        arg_kinds=[ARG_POS, *signature.arg_kinds],
+        arg_names=[None, *signature.arg_names],
+        ret_type=obj_type,
+    )
     return signature
diff --git a/test-data/unit/check-dataclasses.test b/test-data/unit/check-dataclasses.test
@@ -2070,3 +2070,21 @@ class C:
     pass
 
 replace(C())  # E: Argument 1 to "replace" has incompatible type "C"; 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" has incompatible type "str"; expected "int"
+reveal_type(a2)  # N: Revealed type is "__main__.A[builtins.int]"
