Skip to content

Support if statements in dataclass_transform class #14854

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Mar 9, 2023
Merged

Support if statements in dataclass_transform class #14854

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c832d38
Support if statements in dataclass_transform class
KRunchPL Mar 8, 2023
09e77ef
Convert staticmethods into instance methods
KRunchPL Mar 8, 2023
c78dc9c
Use Iterator instead of Generator
KRunchPL Mar 8, 2023
c41f3fb
Additional test when condition is a function
KRunchPL Mar 8, 2023
84daa9b
Adding some elif test cases
KRunchPL Mar 8, 2023
f1b1642
Add Dataclasses test
KRunchPL Mar 8, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 25 additions & 3 deletions mypy/plugins/dataclasses.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from __future__ import annotations

from typing import Optional
from typing import Generator, Optional
from typing_extensions import Final

from mypy import errorcodes, message_registry
Expand All @@ -17,11 +17,13 @@
MDEF,
Argument,
AssignmentStmt,
Block,
CallExpr,
ClassDef,
Context,
DataclassTransformSpec,
Expression,
IfStmt,
JsonDict,
NameExpr,
Node,
Expand Down Expand Up @@ -380,6 +382,26 @@ def reset_init_only_vars(self, info: TypeInfo, attributes: list[DataclassAttribu
# recreate a symbol node for this attribute.
lvalue.node = None

@staticmethod
def _get_assignment_statements_from_if_statement(
stmt: IfStmt,
) -> Generator[AssignmentStmt, None, None]:
for body in stmt.body:
if not body.is_unreachable:
yield from DataclassTransformer._get_assignment_statements_from_block(body)
if stmt.else_body is not None and not stmt.else_body.is_unreachable:
yield from DataclassTransformer._get_assignment_statements_from_block(stmt.else_body)

@staticmethod
def _get_assignment_statements_from_block(
block: Block,
) -> Generator[AssignmentStmt, None, None]:
for stmt in block.body:
if isinstance(stmt, AssignmentStmt):
yield stmt
elif isinstance(stmt, IfStmt):
yield from DataclassTransformer._get_assignment_statements_from_if_statement(stmt)

def collect_attributes(self) -> list[DataclassAttribute] | None:
"""Collect all attributes declared in the dataclass and its parents.

Expand Down Expand Up @@ -438,10 +460,10 @@ def collect_attributes(self) -> list[DataclassAttribute] | None:
# Second, collect attributes belonging to the current class.
current_attr_names: set[str] = set()
kw_only = self._get_bool_arg("kw_only", self._spec.kw_only_default)
for stmt in cls.defs.body:
for stmt in self._get_assignment_statements_from_block(cls.defs):
# Any assignment that doesn't use the new type declaration
# syntax can be ignored out of hand.
if not (isinstance(stmt, AssignmentStmt) and stmt.new_syntax):
if not stmt.new_syntax:
continue

# a: int, b: str = 1, 'foo' is not supported syntax so we
Expand Down
184 changes: 184 additions & 0 deletions test-data/unit/check-dataclass-transform.test
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -451,3 +451,187 @@ Foo(1) # E: Too many arguments for "Foo"

[typing fixtures/typing-full.pyi]
[builtins fixtures/dataclasses.pyi]

[case testDataclassTransformTypeCheckingInFunction]
# flags: --python-version 3.11
from typing import dataclass_transform, Type, TYPE_CHECKING

@dataclass_transform()
def model(cls: Type) -> Type:
return cls

@model
class FunctionModel:
if TYPE_CHECKING:
string_: str
integer_: int
else:
string_: tuple
integer_: tuple

FunctionModel(string_="abc", integer_=1)
FunctionModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "FunctionModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"

[typing fixtures/typing-full.pyi]
[builtins fixtures/dataclasses.pyi]

[case testDataclassTransformNegatedTypeCheckingInFunction]
# flags: --python-version 3.11
from typing import dataclass_transform, Type, TYPE_CHECKING

@dataclass_transform()
def model(cls: Type) -> Type:
return cls

@model
class FunctionModel:
if not TYPE_CHECKING:
string_: tuple
integer_: tuple
else:
string_: str
integer_: int

FunctionModel(string_="abc", integer_=1)
FunctionModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "FunctionModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"

[typing fixtures/typing-full.pyi]
[builtins fixtures/dataclasses.pyi]


[case testDataclassTransformTypeCheckingInBaseClass]
# flags: --python-version 3.11
from typing import dataclass_transform, Type, TYPE_CHECKING

@dataclass_transform()
class ModelBase:
...

class BaseClassModel(ModelBase):
if TYPE_CHECKING:
string_: str
integer_: int
else:
string_: tuple
integer_: tuple

BaseClassModel(string_="abc", integer_=1)
BaseClassModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "BaseClassModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"

[typing fixtures/typing-full.pyi]
[builtins fixtures/dataclasses.pyi]

[case testDataclassTransformNegatedTypeCheckingInBaseClass]
# flags: --python-version 3.11
from typing import dataclass_transform, Type, TYPE_CHECKING

@dataclass_transform()
class ModelBase:
...

class BaseClassModel(ModelBase):
if not TYPE_CHECKING:
string_: tuple
integer_: tuple
else:
string_: str
integer_: int

BaseClassModel(string_="abc", integer_=1)
BaseClassModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "BaseClassModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"

[typing fixtures/typing-full.pyi]
[builtins fixtures/dataclasses.pyi]

[case testDataclassTransformTypeCheckingInMetaClass]
# flags: --python-version 3.11
from typing import dataclass_transform, Type, TYPE_CHECKING

@dataclass_transform()
class ModelMeta(type):
...

class ModelBaseWithMeta(metaclass=ModelMeta):
...

class MetaClassModel(ModelBaseWithMeta):
if TYPE_CHECKING:
string_: str
integer_: int
else:
string_: tuple
integer_: tuple

MetaClassModel(string_="abc", integer_=1)
MetaClassModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "MetaClassModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"

[typing fixtures/typing-full.pyi]
[builtins fixtures/dataclasses.pyi]

[case testDataclassTransformNegatedTypeCheckingInMetaClass]
# flags: --python-version 3.11
from typing import dataclass_transform, Type, TYPE_CHECKING

@dataclass_transform()
class ModelMeta(type):
...

class ModelBaseWithMeta(metaclass=ModelMeta):
...

class MetaClassModel(ModelBaseWithMeta):
if not TYPE_CHECKING:
string_: tuple
integer_: tuple
else:
string_: str
integer_: int

MetaClassModel(string_="abc", integer_=1)
MetaClassModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "MetaClassModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"

[typing fixtures/typing-full.pyi]
[builtins fixtures/dataclasses.pyi]

[case testDataclassTransformConditionalAttributes]
# flags: --python-version 3.11 --always-true TRUTH
from typing import dataclass_transform, Type, TYPE_CHECKING

TRUTH = False # Is set to --always-true

@dataclass_transform()
def model(cls: Type) -> Type:
return cls

@model
class FunctionModel:
if TYPE_CHECKING:
present_1: int
else:
skipped_1: int
if True:
present_2: int
if not False:
present_3: int
if not TRUTH:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's add at least a single test with some def cond() -> bool: ... function.

if cond():
   x: int
   y: int
   z1: int
else:
   x: str
   y: int
   z2: int

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have added the test. It is of course expecting errors, because mypy cannot know what is the return value of a function right? The condition logic is based on IfStmt.is_unreachable which is filled in by logic found in mypy.reachability. I have added a test with some comments that not only tests, but also shows the behavior. I am not sure if we can avoid this limitation (although I am new here, so I can be mistaken). And of course most typical use case will be if TYPE_CHECKING, which we can know is ALWAYS_TRUE.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I think the test is correct 👍

skipped_2: int
else:
present_4: int

FunctionModel(
present_1=1,
present_2=2,
present_3=3,
present_4=4,
)
FunctionModel() # E: Missing positional arguments "present_1", "present_2", "present_3", "present_4" in call to "FunctionModel"
FunctionModel( # E: Unexpected keyword argument "skipped_1" for "FunctionModel"
present_1=1,
present_2=2,
present_3=3,
present_4=4,
skipped_1=5,
)

[typing fixtures/typing-full.pyi]
[builtins fixtures/dataclasses.pyi]