Skip to content

Commit

Permalink
Handle kw_only=True in dataclass fields
Browse files Browse the repository at this point in the history
  • Loading branch information
DanielNoord committed Dec 26, 2022
1 parent a178353 commit 406dfbc
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 14 deletions.
4 changes: 4 additions & 0 deletions ChangeLog
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,10 @@ Release date: TBA

Closes PyCQA/pylint#5225

* Handle the effect of ``kw_only=True`` in dataclass fields correctly.

Closes PyCQA/pylint#7623

* Handle the effect of ``init=False`` in dataclass fields correctly.

Closes PyCQA/pylint#7291
Expand Down
39 changes: 25 additions & 14 deletions astroid/brain/brain_dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,11 +240,12 @@ def _get_previous_field_default(node: nodes.ClassDef, name: str) -> nodes.NodeNG
return None


def _generate_dataclass_init(
def _generate_dataclass_init( # pylint: disable=too-many-locals
node: nodes.ClassDef, assigns: list[nodes.AnnAssign], kw_only_decorated: bool
) -> str:
"""Return an init method for a dataclass given the targets."""
params: list[str] = []
kw_only_params: list[str] = []
assignments: list[str] = []
assign_names: list[str] = []

Expand Down Expand Up @@ -323,7 +324,22 @@ def _generate_dataclass_init(
if previous_default:
param_str += f" = {previous_default.as_string()}"

params.append(param_str)
# If the field is a kw_only field, we need to add it to the kw_only_params
# This overwrites whether or not the class is kw_only decorated
if is_field:
kw_only = [k for k in value.keywords if k.arg == "kw_only"] # type: ignore[union-attr]
if kw_only:
if kw_only[0].value.bool_value():
kw_only_params.append(param_str)
else:
params.append(param_str)
continue
# If kw_only decorated, we need to add all parameters to the kw_only_params
if kw_only_decorated:
kw_only_params.append(param_str)
else:
params.append(param_str)

if not init_var:
assignments.append(assignment_str)

Expand All @@ -332,21 +348,16 @@ def _generate_dataclass_init(
)

# Construct the new init method paramter string
params_string = "self, "
if prev_pos_only:
params_string += prev_pos_only
if not kw_only_decorated:
params_string += ", ".join(params)

# First we do the positional only parameters, making sure to add the
# the self parameter and the comma to allow adding keyword only parameters
params_string = f"self, {prev_pos_only}{', '.join(params)}"
if not params_string.endswith(", "):
params_string += ", "

if prev_kw_only:
params_string += "*, " + prev_kw_only
if kw_only_decorated:
params_string += ", ".join(params) + ", "
elif kw_only_decorated:
params_string += "*, " + ", ".join(params) + ", "
# Then we add the keyword only parameters
if prev_kw_only or kw_only_params:
params_string += "*, "
params_string += f"{prev_kw_only}{', '.join(kw_only_params)}"

assignments_string = "\n ".join(assignments) if assignments else "pass"
return f"def __init__({params_string}) -> None:\n {assignments_string}"
Expand Down
43 changes: 43 additions & 0 deletions tests/unittest_brain_dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -826,6 +826,49 @@ class Dee(Cee):
assert [a.name for a in dee_init.args.kwonlyargs] == []


def test_kw_only_in_field_call() -> None:
"""Test that keyword only fields get correctly put at the end of the __init__."""

first, second, third = astroid.extract_node(
"""
from dataclasses import dataclass, field
@dataclass
class Parent:
p1: int = field(kw_only=True, default=0)
@dataclass
class Child(Parent):
c1: str
@dataclass(kw_only=True)
class GrandChild(Child):
p2: int = field(kw_only=False, default=1)
p3: int = field(kw_only=True, default=2)
Parent.__init__ #@
Child.__init__ #@
GrandChild.__init__ #@
"""
)

first_init: bases.UnboundMethod = next(first.infer())
assert [a.name for a in first_init.args.args] == ["self"]
assert [a.name for a in first_init.args.kwonlyargs] == ["p1"]
assert [d.value for d in first_init.args.kw_defaults] == [0]

second_init: bases.UnboundMethod = next(second.infer())
assert [a.name for a in second_init.args.args] == ["self", "c1"]
assert [a.name for a in second_init.args.kwonlyargs] == ["p1"]
assert [d.value for d in second_init.args.kw_defaults] == [0]

third_init: bases.UnboundMethod = next(third.infer())
assert [a.name for a in third_init.args.args] == ["self", "c1", "p2"]
assert [a.name for a in third_init.args.kwonlyargs] == ["p1", "p3"]
assert [d.value for d in third_init.args.defaults] == [1]
assert [d.value for d in third_init.args.kw_defaults] == [0, 2]


def test_dataclass_with_unknown_base() -> None:
"""Regression test for dataclasses with unknown base classes.
Expand Down

0 comments on commit 406dfbc

Please sign in to comment.