From c832d3832af778b30b98b2bec5116733c33da0d9 Mon Sep 17 00:00:00 2001 From: KRunchPL Date: Wed, 8 Mar 2023 02:53:37 +0100 Subject: [PATCH 1/6] Support if statements in dataclass_transform class Fixes: #14853 --- mypy/plugins/dataclasses.py | 28 ++- test-data/unit/check-dataclass-transform.test | 184 ++++++++++++++++++ 2 files changed, 209 insertions(+), 3 deletions(-) diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py index 7694134ac09e..3274bd9bb651 100644 --- a/mypy/plugins/dataclasses.py +++ b/mypy/plugins/dataclasses.py @@ -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 @@ -17,11 +17,13 @@ MDEF, Argument, AssignmentStmt, + Block, CallExpr, ClassDef, Context, DataclassTransformSpec, Expression, + IfStmt, JsonDict, NameExpr, Node, @@ -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. @@ -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 diff --git a/test-data/unit/check-dataclass-transform.test b/test-data/unit/check-dataclass-transform.test index ec87bd4757ed..cfa3428060e9 100644 --- a/test-data/unit/check-dataclass-transform.test +++ b/test-data/unit/check-dataclass-transform.test @@ -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[, ...]"; 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[, ...]"; 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[, ...]"; 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[, ...]"; 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[, ...]"; 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[, ...]"; 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: + 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] From 09e77ef970d538f90617172027bae75d925dba67 Mon Sep 17 00:00:00 2001 From: KRunchPL Date: Wed, 8 Mar 2023 19:51:50 +0100 Subject: [PATCH 2/6] Convert staticmethods into instance methods --- mypy/plugins/dataclasses.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py index 3274bd9bb651..6227c6567dcf 100644 --- a/mypy/plugins/dataclasses.py +++ b/mypy/plugins/dataclasses.py @@ -382,25 +382,23 @@ 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, + self, 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) + 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 DataclassTransformer._get_assignment_statements_from_block(stmt.else_body) + yield from self._get_assignment_statements_from_block(stmt.else_body) - @staticmethod def _get_assignment_statements_from_block( - block: Block, + self, 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) + 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. From c78dc9c910f10e936b53019b9de836e01f3dd492 Mon Sep 17 00:00:00 2001 From: KRunchPL Date: Wed, 8 Mar 2023 19:58:57 +0100 Subject: [PATCH 3/6] Use Iterator instead of Generator --- mypy/plugins/dataclasses.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py index 6227c6567dcf..a68410765367 100644 --- a/mypy/plugins/dataclasses.py +++ b/mypy/plugins/dataclasses.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Generator, Optional +from typing import Iterator, Optional from typing_extensions import Final from mypy import errorcodes, message_registry @@ -384,16 +384,14 @@ def reset_init_only_vars(self, info: TypeInfo, attributes: list[DataclassAttribu def _get_assignment_statements_from_if_statement( self, stmt: IfStmt - ) -> Generator[AssignmentStmt, None, None]: + ) -> 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 - ) -> Generator[AssignmentStmt, None, None]: + def _get_assignment_statements_from_block(self, block: Block) -> Iterator[AssignmentStmt]: for stmt in block.body: if isinstance(stmt, AssignmentStmt): yield stmt From c41f3fbfa9cb93e7e1b68f158f4cb7d2bd71113e Mon Sep 17 00:00:00 2001 From: KRunchPL Date: Wed, 8 Mar 2023 20:17:48 +0100 Subject: [PATCH 4/6] Additional test when condition is a function --- test-data/unit/check-dataclass-transform.test | 63 +++++++++++++++++-- 1 file changed, 59 insertions(+), 4 deletions(-) diff --git a/test-data/unit/check-dataclass-transform.test b/test-data/unit/check-dataclass-transform.test index cfa3428060e9..d788ab557440 100644 --- a/test-data/unit/check-dataclass-transform.test +++ b/test-data/unit/check-dataclass-transform.test @@ -501,7 +501,7 @@ FunctionModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "Fun [case testDataclassTransformTypeCheckingInBaseClass] # flags: --python-version 3.11 -from typing import dataclass_transform, Type, TYPE_CHECKING +from typing import dataclass_transform, TYPE_CHECKING @dataclass_transform() class ModelBase: @@ -523,7 +523,7 @@ BaseClassModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "Ba [case testDataclassTransformNegatedTypeCheckingInBaseClass] # flags: --python-version 3.11 -from typing import dataclass_transform, Type, TYPE_CHECKING +from typing import dataclass_transform, TYPE_CHECKING @dataclass_transform() class ModelBase: @@ -593,7 +593,7 @@ MetaClassModel(string_="abc", integer_=tuple()) # E: Argument "integer_" to "Me [typing fixtures/typing-full.pyi] [builtins fixtures/dataclasses.pyi] -[case testDataclassTransformConditionalAttributes] +[case testDataclassTransformStaticConditionalAttributes] # flags: --python-version 3.11 --always-true TRUTH from typing import dataclass_transform, Type, TYPE_CHECKING @@ -611,7 +611,7 @@ class FunctionModel: skipped_1: int if True: present_2: int - if not False: + if False: present_3: int if not TRUTH: skipped_2: int @@ -635,3 +635,58 @@ FunctionModel( # E: Unexpected keyword argument "skipped_1" for "FunctionModel [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] From 84daa9bf5b003dacb32c2b2ec886e096067ece68 Mon Sep 17 00:00:00 2001 From: KRunchPL Date: Wed, 8 Mar 2023 20:53:01 +0100 Subject: [PATCH 5/6] Adding some elif test cases --- test-data/unit/check-dataclass-transform.test | 91 ++++++++++++++++++- 1 file changed, 89 insertions(+), 2 deletions(-) diff --git a/test-data/unit/check-dataclass-transform.test b/test-data/unit/check-dataclass-transform.test index d788ab557440..b0c1cdf56097 100644 --- a/test-data/unit/check-dataclass-transform.test +++ b/test-data/unit/check-dataclass-transform.test @@ -609,9 +609,9 @@ class FunctionModel: present_1: int else: skipped_1: int - if True: + if True: # Mypy does not know if it is True or False, so the block is used present_2: int - if False: + if False: # Mypy does not know if it is True or False, so the block is used present_3: int if not TRUTH: skipped_2: int @@ -636,6 +636,93 @@ FunctionModel( # E: Unexpected keyword argument "skipped_1" for "FunctionModel [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 From f1b1642d5c0dbb9e6f2e3f4beca9f0c30fe4bbc4 Mon Sep 17 00:00:00 2001 From: KRunchPL Date: Wed, 8 Mar 2023 20:58:14 +0100 Subject: [PATCH 6/6] Add Dataclasses test --- test-data/unit/check-dataclasses.test | 36 +++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/test-data/unit/check-dataclasses.test b/test-data/unit/check-dataclasses.test index 4d85be391186..da0b7feb4831 100644 --- a/test-data/unit/check-dataclasses.test +++ b/test-data/unit/check-dataclasses.test @@ -2001,3 +2001,39 @@ class Bar(Foo): ... e: Element[Bar] reveal_type(e.elements) # N: Revealed type is "typing.Sequence[__main__.Element[__main__.Bar]]" [builtins fixtures/dataclasses.pyi] + + +[case testIfConditionsInDefinition] +# flags: --python-version 3.11 --always-true TRUTH +from dataclasses import dataclass +from typing import TYPE_CHECKING + +TRUTH = False # Is set to --always-true + +@dataclass +class Foo: + 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: + skipped_2: int + elif 123: + present_4: int + elif TRUTH: + present_5: int + else: + skipped_3: int + +Foo( + present_1=1, + present_2=2, + present_3=3, + present_4=4, + present_5=5, +) +[builtins fixtures/dataclasses.pyi]