Skip to content
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

Jump to bottom

Support if statements in dataclass_transform class #14854

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 all commits
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
24 changes: 21 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 Iterator, 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,22 @@ def reset_init_only_vars(self, info: TypeInfo, attributes: list[DataclassAttribu
# recreate a symbol node for this attribute.
lvalue.node = None

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

def _get_assignment_statements_from_block(self, block: Block) -> Iterator[AssignmentStmt]:
for stmt in block.body:
if isinstance(stmt, AssignmentStmt):
yield stmt
elif isinstance(stmt, IfStmt):
yield from self._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 +456,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
326 changes: 326 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,329 @@ 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_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_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 testDataclassTransformStaticConditionalAttributes]
# 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: # Mypy does not know if it is True or False, so the block is used
present_2: int
if False: # Mypy does not know if it is True or False, so the block is used
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]


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

TRUTH = False # Is set to --always-true
LIE = True # Is set to --always-false

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

@model
class FunctionModel:
if TYPE_CHECKING:
present_1: int
elif TRUTH:
skipped_1: int
else:
skipped_2: int
if LIE:
skipped_3: int
elif TRUTH:
present_2: int
else:
skipped_4: int
if LIE:
skipped_5: int
elif LIE:
skipped_6: int
else:
present_3: int

FunctionModel(
present_1=1,
present_2=2,
present_3=3,
)

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

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

TRUTH = False # Is set to --always-true
LIE = True # Is set to --always-false

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

@model
class FunctionModel:
if 123: # Mypy does not know if it is True or False, so this block is used
present_1: int
elif TRUTH: # Mypy does not know if previous condition is True or False, so it uses also this block
present_2: int
else: # Previous block is for sure True, so this block is skipped
skipped_1: int
if 123:
present_3: int
elif 123:
present_4: int
else:
present_5: int
if 123: # Mypy does not know if it is True or False, so this block is used
present_6: int
elif LIE: # This is for sure False, so the block is skipped used
skipped_2: int
else: # None of the conditions above for sure True, so this block is used
present_7: int

FunctionModel(
present_1=1,
present_2=2,
present_3=3,
present_4=4,
present_5=5,
present_6=6,
present_7=7,
)

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

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

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

def condition() -> bool:
return True

@model
class FunctionModel:
if condition():
x: int
y: int
z1: int
else:
x: str # E: Name "x" already defined on line 14
y: int # E: Name "y" already defined on line 15
z2: int

FunctionModel(x=1, y=2, z1=3, z2=4)

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


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

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

def condition() -> bool:
return True

@model
class FunctionModel:
if not condition():
x: int
y: int
z1: int
else:
x: str # E: Name "x" already defined on line 14
y: int # E: Name "y" already defined on line 15
z2: int

FunctionModel(x=1, y=2, z1=3, z2=4)

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