Skip to content

Commit 2e8dcbf

Browse files
KRunchPLKRunchPL
authored
Support if statements in dataclass and dataclass_transform plugin (#14854)
Fixes: #14853 Adding support for `if` statements in the `dataclass` and `dataclass_transform` decorators, so the attributes defined conditionally are treated as those that are directly in class body.
1 parent dfe0281 commit 2e8dcbf

File tree

3 files changed

+383
-3
lines changed

3 files changed

+383
-3
lines changed

‎mypy/plugins/dataclasses.py

+21-3
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from __future__ import annotations
44

5-
from typing import Optional
5+
from typing import Iterator, Optional
66
from typing_extensions import Final
77

88
from mypy import errorcodes, message_registry
@@ -17,11 +17,13 @@
1717
MDEF,
1818
Argument,
1919
AssignmentStmt,
20+
Block,
2021
CallExpr,
2122
ClassDef,
2223
Context,
2324
DataclassTransformSpec,
2425
Expression,
26+
IfStmt,
2527
JsonDict,
2628
NameExpr,
2729
Node,
@@ -380,6 +382,22 @@ def reset_init_only_vars(self, info: TypeInfo, attributes: list[DataclassAttribu
380382
# recreate a symbol node for this attribute.
381383
lvalue.node = None
382384

385+
def _get_assignment_statements_from_if_statement(
386+
self, stmt: IfStmt
387+
) -> Iterator[AssignmentStmt]:
388+
for body in stmt.body:
389+
if not body.is_unreachable:
390+
yield from self._get_assignment_statements_from_block(body)
391+
if stmt.else_body is not None and not stmt.else_body.is_unreachable:
392+
yield from self._get_assignment_statements_from_block(stmt.else_body)
393+
394+
def _get_assignment_statements_from_block(self, block: Block) -> Iterator[AssignmentStmt]:
395+
for stmt in block.body:
396+
if isinstance(stmt, AssignmentStmt):
397+
yield stmt
398+
elif isinstance(stmt, IfStmt):
399+
yield from self._get_assignment_statements_from_if_statement(stmt)
400+
383401
def collect_attributes(self) -> list[DataclassAttribute] | None:
384402
"""Collect all attributes declared in the dataclass and its parents.
385403
@@ -438,10 +456,10 @@ def collect_attributes(self) -> list[DataclassAttribute] | None:
438456
# Second, collect attributes belonging to the current class.
439457
current_attr_names: set[str] = set()
440458
kw_only = self._get_bool_arg("kw_only", self._spec.kw_only_default)
441-
for stmt in cls.defs.body:
459+
for stmt in self._get_assignment_statements_from_block(cls.defs):
442460
# Any assignment that doesn't use the new type declaration
443461
# syntax can be ignored out of hand.
444-
if not (isinstance(stmt, AssignmentStmt) and stmt.new_syntax):
462+
if not stmt.new_syntax:
445463
continue
446464

447465
# a: int, b: str = 1, 'foo' is not supported syntax so we

‎test-data/unit/check-dataclass-transform.test

+326
Original file line numberDiff line numberDiff line change
@@ -451,3 +451,329 @@ Foo(1) # E: Too many arguments for "Foo"
451451

452452
[typing fixtures/typing-full.pyi]
453453
[builtins fixtures/dataclasses.pyi]
454+
455+
[case testDataclassTransformTypeCheckingInFunction]
456+
# flags: --python-version 3.11
457+
from typing import dataclass_transform, Type, TYPE_CHECKING
458+
459+
@dataclass_transform()
460+
def model(cls: Type) -> Type:
461+
return cls
462+
463+
@model
464+
class FunctionModel:
465+
if TYPE_CHECKING:
466+
string_: str
467+
integer_: int
468+
else:
469+
string_: tuple
470+
integer_: tuple
471+
472+
FunctionModel(string_="abc", integer_=1)
473+
FunctionModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "FunctionModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"
474+
475+
[typing fixtures/typing-full.pyi]
476+
[builtins fixtures/dataclasses.pyi]
477+
478+
[case testDataclassTransformNegatedTypeCheckingInFunction]
479+
# flags: --python-version 3.11
480+
from typing import dataclass_transform, Type, TYPE_CHECKING
481+
482+
@dataclass_transform()
483+
def model(cls: Type) -> Type:
484+
return cls
485+
486+
@model
487+
class FunctionModel:
488+
if not TYPE_CHECKING:
489+
string_: tuple
490+
integer_: tuple
491+
else:
492+
string_: str
493+
integer_: int
494+
495+
FunctionModel(string_="abc", integer_=1)
496+
FunctionModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "FunctionModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"
497+
498+
[typing fixtures/typing-full.pyi]
499+
[builtins fixtures/dataclasses.pyi]
500+
501+
502+
[case testDataclassTransformTypeCheckingInBaseClass]
503+
# flags: --python-version 3.11
504+
from typing import dataclass_transform, TYPE_CHECKING
505+
506+
@dataclass_transform()
507+
class ModelBase:
508+
...
509+
510+
class BaseClassModel(ModelBase):
511+
if TYPE_CHECKING:
512+
string_: str
513+
integer_: int
514+
else:
515+
string_: tuple
516+
integer_: tuple
517+
518+
BaseClassModel(string_="abc", integer_=1)
519+
BaseClassModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "BaseClassModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"
520+
521+
[typing fixtures/typing-full.pyi]
522+
[builtins fixtures/dataclasses.pyi]
523+
524+
[case testDataclassTransformNegatedTypeCheckingInBaseClass]
525+
# flags: --python-version 3.11
526+
from typing import dataclass_transform, TYPE_CHECKING
527+
528+
@dataclass_transform()
529+
class ModelBase:
530+
...
531+
532+
class BaseClassModel(ModelBase):
533+
if not TYPE_CHECKING:
534+
string_: tuple
535+
integer_: tuple
536+
else:
537+
string_: str
538+
integer_: int
539+
540+
BaseClassModel(string_="abc", integer_=1)
541+
BaseClassModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "BaseClassModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"
542+
543+
[typing fixtures/typing-full.pyi]
544+
[builtins fixtures/dataclasses.pyi]
545+
546+
[case testDataclassTransformTypeCheckingInMetaClass]
547+
# flags: --python-version 3.11
548+
from typing import dataclass_transform, Type, TYPE_CHECKING
549+
550+
@dataclass_transform()
551+
class ModelMeta(type):
552+
...
553+
554+
class ModelBaseWithMeta(metaclass=ModelMeta):
555+
...
556+
557+
class MetaClassModel(ModelBaseWithMeta):
558+
if TYPE_CHECKING:
559+
string_: str
560+
integer_: int
561+
else:
562+
string_: tuple
563+
integer_: tuple
564+
565+
MetaClassModel(string_="abc", integer_=1)
566+
MetaClassModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "MetaClassModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"
567+
568+
[typing fixtures/typing-full.pyi]
569+
[builtins fixtures/dataclasses.pyi]
570+
571+
[case testDataclassTransformNegatedTypeCheckingInMetaClass]
572+
# flags: --python-version 3.11
573+
from typing import dataclass_transform, Type, TYPE_CHECKING
574+
575+
@dataclass_transform()
576+
class ModelMeta(type):
577+
...
578+
579+
class ModelBaseWithMeta(metaclass=ModelMeta):
580+
...
581+
582+
class MetaClassModel(ModelBaseWithMeta):
583+
if not TYPE_CHECKING:
584+
string_: tuple
585+
integer_: tuple
586+
else:
587+
string_: str
588+
integer_: int
589+
590+
MetaClassModel(string_="abc", integer_=1)
591+
MetaClassModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "MetaClassModel" has incompatible type "Tuple[<nothing>, ...]"; expected "int"
592+
593+
[typing fixtures/typing-full.pyi]
594+
[builtins fixtures/dataclasses.pyi]
595+
596+
[case testDataclassTransformStaticConditionalAttributes]
597+
# flags: --python-version 3.11 --always-true TRUTH
598+
from typing import dataclass_transform, Type, TYPE_CHECKING
599+
600+
TRUTH = False # Is set to --always-true
601+
602+
@dataclass_transform()
603+
def model(cls: Type) -> Type:
604+
return cls
605+
606+
@model
607+
class FunctionModel:
608+
if TYPE_CHECKING:
609+
present_1: int
610+
else:
611+
skipped_1: int
612+
if True: # Mypy does not know if it is True or False, so the block is used
613+
present_2: int
614+
if False: # Mypy does not know if it is True or False, so the block is used
615+
present_3: int
616+
if not TRUTH:
617+
skipped_2: int
618+
else:
619+
present_4: int
620+
621+
FunctionModel(
622+
present_1=1,
623+
present_2=2,
624+
present_3=3,
625+
present_4=4,
626+
)
627+
FunctionModel() # E: Missing positional arguments "present_1", "present_2", "present_3", "present_4" in call to "FunctionModel"
628+
FunctionModel( # E: Unexpected keyword argument "skipped_1" for "FunctionModel"
629+
present_1=1,
630+
present_2=2,
631+
present_3=3,
632+
present_4=4,
633+
skipped_1=5,
634+
)
635+
636+
[typing fixtures/typing-full.pyi]
637+
[builtins fixtures/dataclasses.pyi]
638+
639+
640+
[case testDataclassTransformStaticDeterministicConditionalElifAttributes]
641+
# flags: --python-version 3.11 --always-true TRUTH --always-false LIE
642+
from typing import dataclass_transform, Type, TYPE_CHECKING
643+
644+
TRUTH = False # Is set to --always-true
645+
LIE = True # Is set to --always-false
646+
647+
@dataclass_transform()
648+
def model(cls: Type) -> Type:
649+
return cls
650+
651+
@model
652+
class FunctionModel:
653+
if TYPE_CHECKING:
654+
present_1: int
655+
elif TRUTH:
656+
skipped_1: int
657+
else:
658+
skipped_2: int
659+
if LIE:
660+
skipped_3: int
661+
elif TRUTH:
662+
present_2: int
663+
else:
664+
skipped_4: int
665+
if LIE:
666+
skipped_5: int
667+
elif LIE:
668+
skipped_6: int
669+
else:
670+
present_3: int
671+
672+
FunctionModel(
673+
present_1=1,
674+
present_2=2,
675+
present_3=3,
676+
)
677+
678+
[typing fixtures/typing-full.pyi]
679+
[builtins fixtures/dataclasses.pyi]
680+
681+
[case testDataclassTransformStaticNotDeterministicConditionalElifAttributes]
682+
# flags: --python-version 3.11 --always-true TRUTH --always-false LIE
683+
from typing import dataclass_transform, Type, TYPE_CHECKING
684+
685+
TRUTH = False # Is set to --always-true
686+
LIE = True # Is set to --always-false
687+
688+
@dataclass_transform()
689+
def model(cls: Type) -> Type:
690+
return cls
691+
692+
@model
693+
class FunctionModel:
694+
if 123: # Mypy does not know if it is True or False, so this block is used
695+
present_1: int
696+
elif TRUTH: # Mypy does not know if previous condition is True or False, so it uses also this block
697+
present_2: int
698+
else: # Previous block is for sure True, so this block is skipped
699+
skipped_1: int
700+
if 123:
701+
present_3: int
702+
elif 123:
703+
present_4: int
704+
else:
705+
present_5: int
706+
if 123: # Mypy does not know if it is True or False, so this block is used
707+
present_6: int
708+
elif LIE: # This is for sure False, so the block is skipped used
709+
skipped_2: int
710+
else: # None of the conditions above for sure True, so this block is used
711+
present_7: int
712+
713+
FunctionModel(
714+
present_1=1,
715+
present_2=2,
716+
present_3=3,
717+
present_4=4,
718+
present_5=5,
719+
present_6=6,
720+
present_7=7,
721+
)
722+
723+
[typing fixtures/typing-full.pyi]
724+
[builtins fixtures/dataclasses.pyi]
725+
726+
[case testDataclassTransformFunctionConditionalAttributes]
727+
# flags: --python-version 3.11
728+
from typing import dataclass_transform, Type
729+
730+
@dataclass_transform()
731+
def model(cls: Type) -> Type:
732+
return cls
733+
734+
def condition() -> bool:
735+
return True
736+
737+
@model
738+
class FunctionModel:
739+
if condition():
740+
x: int
741+
y: int
742+
z1: int
743+
else:
744+
x: str # E: Name "x" already defined on line 14
745+
y: int # E: Name "y" already defined on line 15
746+
z2: int
747+
748+
FunctionModel(x=1, y=2, z1=3, z2=4)
749+
750+
[typing fixtures/typing-full.pyi]
751+
[builtins fixtures/dataclasses.pyi]
752+
753+
754+
[case testDataclassTransformNegatedFunctionConditionalAttributes]
755+
# flags: --python-version 3.11
756+
from typing import dataclass_transform, Type
757+
758+
@dataclass_transform()
759+
def model(cls: Type) -> Type:
760+
return cls
761+
762+
def condition() -> bool:
763+
return True
764+
765+
@model
766+
class FunctionModel:
767+
if not condition():
768+
x: int
769+
y: int
770+
z1: int
771+
else:
772+
x: str # E: Name "x" already defined on line 14
773+
y: int # E: Name "y" already defined on line 15
774+
z2: int
775+
776+
FunctionModel(x=1, y=2, z1=3, z2=4)
777+
778+
[typing fixtures/typing-full.pyi]
779+
[builtins fixtures/dataclasses.pyi]

0 commit comments

Comments
 (0)