Skip to content

Commit dbfd425

Browse files
committed
[dataclasses plugin] Support kw_only=True
1 parent 97b3b90 commit dbfd425

File tree

3 files changed

+124
-20
lines changed

3 files changed

+124
-20
lines changed

mypy/plugins/dataclasses.py

+19-2
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ def __init__(
3939
column: int,
4040
type: Optional[Type],
4141
info: TypeInfo,
42+
kw_only: bool,
4243
) -> None:
4344
self.name = name
4445
self.is_in_init = is_in_init
@@ -48,6 +49,7 @@ def __init__(
4849
self.column = column
4950
self.type = type
5051
self.info = info
52+
self.kw_only = kw_only
5153

5254
def to_argument(self) -> Argument:
5355
return Argument(
@@ -77,6 +79,8 @@ def deserialize(
7779
cls, info: TypeInfo, data: JsonDict, api: SemanticAnalyzerPluginInterface
7880
) -> 'DataclassAttribute':
7981
data = data.copy()
82+
if data.get('kw_only') is None:
83+
data['kw_only'] = False
8084
typ = deserialize_and_fixup_type(data.pop('type'), api)
8185
return cls(type=typ, info=info, **data)
8286

@@ -215,6 +219,7 @@ def collect_attributes(self) -> Optional[List[DataclassAttribute]]:
215219
cls = self._ctx.cls
216220
attrs: List[DataclassAttribute] = []
217221
known_attrs: Set[str] = set()
222+
kw_only = _get_decorator_bool_argument(ctx, 'kw_only', False)
218223
for stmt in cls.defs.body:
219224
# Any assignment that doesn't use the new type declaration
220225
# syntax can be ignored out of hand.
@@ -251,6 +256,10 @@ def collect_attributes(self) -> Optional[List[DataclassAttribute]]:
251256
is_init_var = True
252257
node.type = node_type.args[0]
253258

259+
if (isinstance(node_type, Instance) and
260+
node_type.type.fullname == 'dataclasses._KW_ONLY_TYPE'):
261+
kw_only = True
262+
254263
has_field_call, field_args = _collect_field_args(stmt.rvalue)
255264

256265
is_in_init_param = field_args.get('init')
@@ -274,6 +283,13 @@ def collect_attributes(self) -> Optional[List[DataclassAttribute]]:
274283
# on self in the generated __init__(), not in the class body.
275284
sym.implicit = True
276285

286+
is_kw_only = kw_only
287+
# Use the kw_only field arg if it is provided. Otherwise use the
288+
# kw_only value from the decorator parameter.
289+
field_kw_only_param = field_args.get('kw_only')
290+
if field_kw_only_param is not None:
291+
is_kw_only = bool(ctx.api.parse_bool(field_kw_only_param))
292+
277293
known_attrs.add(lhs.name)
278294
attrs.append(DataclassAttribute(
279295
name=lhs.name,
@@ -284,6 +300,7 @@ def collect_attributes(self) -> Optional[List[DataclassAttribute]]:
284300
column=stmt.column,
285301
type=sym.type,
286302
info=cls.info,
303+
kw_only=is_kw_only,
287304
))
288305

289306
# Next, collect attributes belonging to any class in the MRO
@@ -323,10 +340,10 @@ def collect_attributes(self) -> Optional[List[DataclassAttribute]]:
323340
# arguments that have a default.
324341
found_default = False
325342
for attr in all_attrs:
326-
# If we find any attribute that is_in_init but that
343+
# If we find any attribute that is_in_init, not kw_only, and that
327344
# doesn't have a default after one that does have one,
328345
# then that's an error.
329-
if found_default and attr.is_in_init and not attr.has_default:
346+
if found_default and attr.is_in_init and not attr.has_default and not attr.kw_only:
330347
# If the issue comes from merging different classes, report it
331348
# at the class definition point.
332349
context = (Context(line=attr.line, column=attr.column) if attr in attrs

test-data/unit/check-dataclasses.test

+58-1
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ class Person:
225225
name: str
226226
age: int = field(init=None) # E: No overload variant of "field" matches argument type "None" \
227227
# N: Possible overload variant: \
228-
# N: def field(*, init: bool = ..., repr: bool = ..., hash: Optional[bool] = ..., compare: bool = ..., metadata: Optional[Mapping[str, Any]] = ...) -> Any \
228+
# N: def field(*, init: bool = ..., repr: bool = ..., hash: Optional[bool] = ..., compare: bool = ..., metadata: Optional[Mapping[str, Any]] = ..., kw_only: bool = ...) -> Any \
229229
# N: <2 more non-matching overloads not shown>
230230

231231
[builtins fixtures/list.pyi]
@@ -311,6 +311,63 @@ class Application:
311311

312312
[builtins fixtures/list.pyi]
313313

314+
[case testDataclassesOrderingKwOnly]
315+
# flags: --python-version 3.10
316+
from dataclasses import dataclass
317+
318+
@dataclass(kw_only=True)
319+
class Application:
320+
name: str = 'Unnamed'
321+
rating: int
322+
323+
[builtins fixtures/list.pyi]
324+
325+
[case testDataclassesOrderingKwOnlyOnField]
326+
# flags: --python-version 3.10
327+
from dataclasses import dataclass, field
328+
329+
@dataclass
330+
class Application:
331+
name: str = 'Unnamed'
332+
rating: int = field(kw_only=True)
333+
334+
[builtins fixtures/list.pyi]
335+
336+
[case testDataclassesOrderingKwOnlyOnFieldFalse]
337+
# flags: --python-version 3.10
338+
from dataclasses import dataclass, field
339+
340+
@dataclass
341+
class Application:
342+
name: str = 'Unnamed'
343+
rating: int = field(kw_only=False) # E: Attributes without a default cannot follow attributes with one
344+
345+
[builtins fixtures/list.pyi]
346+
347+
[case testDataclassesOrderingKwOnlyWithSentinel]
348+
# flags: --python-version 3.10
349+
from dataclasses import dataclass, KW_ONLY
350+
351+
@dataclass
352+
class Application:
353+
_: KW_ONLY
354+
name: str = 'Unnamed'
355+
rating: int
356+
357+
[builtins fixtures/list.pyi]
358+
359+
[case testDataclassesOrderingKwOnlyWithSentinelAndFieldOverride]
360+
# flags: --python-version 3.10
361+
from dataclasses import dataclass, field, KW_ONLY
362+
363+
@dataclass
364+
class Application:
365+
_: KW_ONLY
366+
name: str = 'Unnamed'
367+
rating: int = field(kw_only=False) # E: Attributes without a default cannot follow attributes with one
368+
369+
[builtins fixtures/list.pyi]
370+
314371
[case testDataclassesClassmethods]
315372
# flags: --python-version 3.7
316373
from dataclasses import dataclass

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

+47-17
Original file line numberDiff line numberDiff line change
@@ -5,26 +5,56 @@ _T = TypeVar('_T')
55
class InitVar(Generic[_T]):
66
...
77

8+
class _KW_ONLY_TYPE: ...
9+
KW_ONLY = _KW_ONLY_TYPE
810

911
@overload
10-
def dataclass(_cls: Type[_T]) -> Type[_T]: ...
11-
12+
def dataclass(__cls: Type[_T]) -> Type[_T]: ...
1213
@overload
13-
def dataclass(*, init: bool = ..., repr: bool = ..., eq: bool = ..., order: bool = ...,
14-
unsafe_hash: bool = ..., frozen: bool = ...) -> Callable[[Type[_T]], Type[_T]]: ...
15-
16-
14+
def dataclass(__cls: None) -> Callable[[Type[_T]], Type[_T]]: ...
1715
@overload
18-
def field(*, default: _T,
19-
init: bool = ..., repr: bool = ..., hash: Optional[bool] = ..., compare: bool = ...,
20-
metadata: Optional[Mapping[str, Any]] = ...) -> _T: ...
21-
16+
def dataclass(
17+
*,
18+
init: bool = ...,
19+
repr: bool = ...,
20+
eq: bool = ...,
21+
order: bool = ...,
22+
unsafe_hash: bool = ...,
23+
frozen: bool = ...,
24+
match_args: bool = ...,
25+
kw_only: bool = ...,
26+
slots: bool = ...,
27+
) -> Callable[[Type[_T]], Type[_T]]: ...
28+
29+
@overload # `default` and `default_factory` are optional and mutually exclusive.
30+
def field(
31+
*,
32+
default: _T,
33+
init: bool = ...,
34+
repr: bool = ...,
35+
hash: Optional[bool] = ...,
36+
compare: bool = ...,
37+
metadata: Optional[Mapping[str, Any]] = ...,
38+
kw_only: bool = ...,
39+
) -> _T: ...
2240
@overload
23-
def field(*, default_factory: Callable[[], _T],
24-
init: bool = ..., repr: bool = ..., hash: Optional[bool] = ..., compare: bool = ...,
25-
metadata: Optional[Mapping[str, Any]] = ...) -> _T: ...
26-
41+
def field(
42+
*,
43+
default_factory: Callable[[], _T],
44+
init: bool = ...,
45+
repr: bool = ...,
46+
hash: Optional[bool] = ...,
47+
compare: bool = ...,
48+
metadata: Optional[Mapping[str, Any]] = ...,
49+
kw_only: bool = ...,
50+
) -> _T: ...
2751
@overload
28-
def field(*,
29-
init: bool = ..., repr: bool = ..., hash: Optional[bool] = ..., compare: bool = ...,
30-
metadata: Optional[Mapping[str, Any]] = ...) -> Any: ...
52+
def field(
53+
*,
54+
init: bool = ...,
55+
repr: bool = ...,
56+
hash: Optional[bool] = ...,
57+
compare: bool = ...,
58+
metadata: Optional[Mapping[str, Any]] = ...,
59+
kw_only: bool = ...,
60+
) -> Any: ...

0 commit comments

Comments
 (0)