|
2 | 2 |
|
3 | 3 | from __future__ import annotations |
4 | 4 |
|
5 | | -from typing import TYPE_CHECKING, Final, Iterator |
| 5 | +from typing import TYPE_CHECKING, Final, Iterator, Literal |
6 | 6 |
|
7 | 7 | from mypy import errorcodes, message_registry |
8 | 8 | from mypy.expandtype import expand_type, expand_type_by_instance |
|
86 | 86 | field_specifiers=("dataclasses.Field", "dataclasses.field"), |
87 | 87 | ) |
88 | 88 | _INTERNAL_REPLACE_SYM_NAME: Final = "__mypy-replace" |
89 | | -_INTERNAL_POST_INIT_SYM_NAME: Final = "__mypy-__post_init__" |
| 89 | +_INTERNAL_POST_INIT_SYM_NAME: Final = "__mypy-post_init" |
90 | 90 |
|
91 | 91 |
|
92 | 92 | class DataclassAttribute: |
@@ -118,14 +118,33 @@ def __init__( |
118 | 118 | self.is_neither_frozen_nor_nonfrozen = is_neither_frozen_nor_nonfrozen |
119 | 119 | self._api = api |
120 | 120 |
|
121 | | - def to_argument(self, current_info: TypeInfo) -> Argument: |
122 | | - arg_kind = ARG_POS |
123 | | - if self.kw_only and self.has_default: |
124 | | - arg_kind = ARG_NAMED_OPT |
125 | | - elif self.kw_only and not self.has_default: |
126 | | - arg_kind = ARG_NAMED |
127 | | - elif not self.kw_only and self.has_default: |
128 | | - arg_kind = ARG_OPT |
| 121 | + def to_argument( |
| 122 | + self, current_info: TypeInfo, *, of: Literal["__init__", "replace", "__post_init__"] |
| 123 | + ) -> Argument: |
| 124 | + if of == "__init__": |
| 125 | + arg_kind = ARG_POS |
| 126 | + if self.kw_only and self.has_default: |
| 127 | + arg_kind = ARG_NAMED_OPT |
| 128 | + elif self.kw_only and not self.has_default: |
| 129 | + arg_kind = ARG_NAMED |
| 130 | + elif not self.kw_only and self.has_default: |
| 131 | + arg_kind = ARG_OPT |
| 132 | + elif of == "replace": |
| 133 | + arg_kind = ARG_NAMED if self.is_init_var and not self.has_default else ARG_NAMED_OPT |
| 134 | + elif of == "__post_init__": |
| 135 | + # We always use `ARG_POS` without a default value, because it is practical. |
| 136 | + # Consider this case: |
| 137 | + # |
| 138 | + # @dataclass |
| 139 | + # class My: |
| 140 | + # y: dataclasses.InitVar[str] = 'a' |
| 141 | + # def __post_init__(self, y: str) -> None: ... |
| 142 | + # |
| 143 | + # We would be *required* to specify `y: str = ...` if default is added here. |
| 144 | + # But, most people won't care about adding default values to `__post_init__`, |
| 145 | + # because it is not designed to be called directly, and duplicating default values |
| 146 | + # for the sake of type-checking is unpleasant. |
| 147 | + arg_kind = ARG_POS |
129 | 148 | return Argument( |
130 | 149 | variable=self.to_var(current_info), |
131 | 150 | type_annotation=self.expand_type(current_info), |
@@ -236,7 +255,7 @@ def transform(self) -> bool: |
236 | 255 | and attributes |
237 | 256 | ): |
238 | 257 | args = [ |
239 | | - attr.to_argument(info) |
| 258 | + attr.to_argument(info, of="__init__") |
240 | 259 | for attr in attributes |
241 | 260 | if attr.is_in_init and not self._is_kw_only_type(attr.type) |
242 | 261 | ] |
@@ -375,70 +394,26 @@ def _add_internal_replace_method(self, attributes: list[DataclassAttribute]) -> |
375 | 394 | Stashes the signature of 'dataclasses.replace(...)' for this specific dataclass |
376 | 395 | to be used later whenever 'dataclasses.replace' is called for this dataclass. |
377 | 396 | """ |
378 | | - arg_types: list[Type] = [] |
379 | | - arg_kinds = [] |
380 | | - arg_names: list[str | None] = [] |
381 | | - |
382 | | - info = self._cls.info |
383 | | - for attr in attributes: |
384 | | - attr_type = attr.expand_type(info) |
385 | | - assert attr_type is not None |
386 | | - arg_types.append(attr_type) |
387 | | - arg_kinds.append( |
388 | | - ARG_NAMED if attr.is_init_var and not attr.has_default else ARG_NAMED_OPT |
389 | | - ) |
390 | | - arg_names.append(attr.name) |
391 | | - |
392 | | - signature = CallableType( |
393 | | - arg_types=arg_types, |
394 | | - arg_kinds=arg_kinds, |
395 | | - arg_names=arg_names, |
396 | | - ret_type=NoneType(), |
397 | | - fallback=self._api.named_type("builtins.function"), |
398 | | - ) |
399 | | - |
400 | | - info.names[_INTERNAL_REPLACE_SYM_NAME] = SymbolTableNode( |
401 | | - kind=MDEF, node=FuncDef(typ=signature), plugin_generated=True |
| 397 | + add_method_to_class( |
| 398 | + self._api, |
| 399 | + self._cls, |
| 400 | + _INTERNAL_REPLACE_SYM_NAME, |
| 401 | + args=[attr.to_argument(self._cls.info, of="replace") for attr in attributes], |
| 402 | + return_type=NoneType(), |
| 403 | + is_staticmethod=True, |
402 | 404 | ) |
403 | 405 |
|
404 | 406 | def _add_internal_post_init_method(self, attributes: list[DataclassAttribute]) -> None: |
405 | | - arg_types: list[Type] = [fill_typevars(self._cls.info)] |
406 | | - arg_kinds = [ARG_POS] |
407 | | - arg_names: list[str | None] = ["self"] |
408 | | - |
409 | | - info = self._cls.info |
410 | | - for attr in attributes: |
411 | | - if not attr.is_init_var: |
412 | | - continue |
413 | | - attr_type = attr.expand_type(info) |
414 | | - assert attr_type is not None |
415 | | - arg_types.append(attr_type) |
416 | | - # We always use `ARG_POS` without a default value, because it is practical. |
417 | | - # Consider this case: |
418 | | - # |
419 | | - # @dataclass |
420 | | - # class My: |
421 | | - # y: dataclasses.InitVar[str] = 'a' |
422 | | - # def __post_init__(self, y: str) -> None: ... |
423 | | - # |
424 | | - # We would be *required* to specify `y: str = ...` if default is added here. |
425 | | - # But, most people won't care about adding default values to `__post_init__`, |
426 | | - # because it is not designed to be called directly, and duplicating default values |
427 | | - # for the sake of type-checking is unpleasant. |
428 | | - arg_kinds.append(ARG_POS) |
429 | | - arg_names.append(attr.name) |
430 | | - |
431 | | - signature = CallableType( |
432 | | - arg_types=arg_types, |
433 | | - arg_kinds=arg_kinds, |
434 | | - arg_names=arg_names, |
435 | | - ret_type=NoneType(), |
436 | | - fallback=self._api.named_type("builtins.function"), |
437 | | - name="__post_init__", |
438 | | - ) |
439 | | - |
440 | | - info.names[_INTERNAL_POST_INIT_SYM_NAME] = SymbolTableNode( |
441 | | - kind=MDEF, node=FuncDef(typ=signature), plugin_generated=True |
| 407 | + add_method_to_class( |
| 408 | + self._api, |
| 409 | + self._cls, |
| 410 | + _INTERNAL_POST_INIT_SYM_NAME, |
| 411 | + args=[ |
| 412 | + attr.to_argument(self._cls.info, of="__post_init__") |
| 413 | + for attr in attributes |
| 414 | + if attr.is_init_var |
| 415 | + ], |
| 416 | + return_type=NoneType(), |
442 | 417 | ) |
443 | 418 |
|
444 | 419 | def add_slots( |
@@ -1120,20 +1095,18 @@ def is_processed_dataclass(info: TypeInfo | None) -> bool: |
1120 | 1095 | def check_post_init(api: TypeChecker, defn: FuncItem, info: TypeInfo) -> None: |
1121 | 1096 | if defn.type is None: |
1122 | 1097 | return |
1123 | | - |
1124 | | - ideal_sig = info.get_method(_INTERNAL_POST_INIT_SYM_NAME) |
1125 | | - if ideal_sig is None or ideal_sig.type is None: |
1126 | | - return |
1127 | | - |
1128 | | - # We set it ourself, so it is always fine: |
1129 | | - assert isinstance(ideal_sig.type, ProperType) |
1130 | | - assert isinstance(ideal_sig.type, FunctionLike) |
1131 | | - # Type of `FuncItem` is always `FunctionLike`: |
1132 | 1098 | assert isinstance(defn.type, FunctionLike) |
1133 | 1099 |
|
| 1100 | + ideal_sig_method = info.get_method(_INTERNAL_POST_INIT_SYM_NAME) |
| 1101 | + assert ideal_sig_method is not None and ideal_sig_method.type is not None |
| 1102 | + ideal_sig = ideal_sig_method.type |
| 1103 | + assert isinstance(ideal_sig, ProperType) # we set it ourselves |
| 1104 | + assert isinstance(ideal_sig, CallableType) |
| 1105 | + ideal_sig = ideal_sig.copy_modified(name="__post_init__") |
| 1106 | + |
1134 | 1107 | api.check_override( |
1135 | 1108 | override=defn.type, |
1136 | | - original=ideal_sig.type, |
| 1109 | + original=ideal_sig, |
1137 | 1110 | name="__post_init__", |
1138 | 1111 | name_in_super="__post_init__", |
1139 | 1112 | supertype="dataclass", |
|
0 commit comments