Skip to content

Commit

Permalink
Generic dataclass support
Browse files Browse the repository at this point in the history
  • Loading branch information
ikonst committed Mar 30, 2023
1 parent 40315b7 commit c005895
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 7 deletions.
25 changes: 18 additions & 7 deletions mypy/plugins/dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
)
Expand All @@ -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}",
)
Expand Down Expand Up @@ -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
18 changes: 18 additions & 0 deletions test-data/unit/check-dataclasses.test
Original file line number Diff line number Diff line change
Expand Up @@ -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]"

0 comments on commit c005895

Please sign in to comment.