|
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