Use correct type for AugAssign and AnnAssign target (#396)

* Use correct type * Add tests * Suppress intentional type errors in pyre
Instagram · Oct 1, 2020 · 6731aa5 · 6731aa5
1 parent 10d6451
commit 6731aa5
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 18 deletions.
diff --git a/libcst/_nodes/statement.py b/libcst/_nodes/statement.py
@@ -1332,7 +1332,7 @@ class AnnAssign(BaseSmallStatement):
     """
 
     #: The target that is being annotated and possibly assigned to.
-    target: BaseExpression
+    target: BaseAssignTargetExpression
 
     #: The annotation for the target.
     annotation: Annotation
@@ -1393,7 +1393,7 @@ class AugAssign(BaseSmallStatement):
     """
 
     #: Target that is being operated on and assigned to.
-    target: BaseExpression
+    target: BaseAssignTargetExpression
 
     #: The augmented assignment operation being performed.
     operator: BaseAugOp

diff --git a/libcst/_nodes/tests/base.py b/libcst/_nodes/tests/base.py
@@ -95,6 +95,12 @@ def assert_invalid(
         with self.assertRaisesRegex(cst.CSTValidationError, expected_re):
             get_node()
 
+    def assert_invalid_types(
+        self, get_node: Callable[[], cst.CSTNode], expected_re: str
+    ) -> None:
+        with self.assertRaisesRegex(TypeError, expected_re):
+            get_node().validate_types_shallow()
+
     def __assert_codegen(
         self,
         node: cst.CSTNode,

diff --git a/libcst/_nodes/tests/test_assign.py b/libcst/_nodes/tests/test_assign.py
@@ -119,6 +119,29 @@ def test_valid(self, **kwargs: Any) -> None:
     def test_invalid(self, **kwargs: Any) -> None:
         self.assert_invalid(**kwargs)
 
+    @data_provider(
+        (
+            {
+                "get_node": (
+                    lambda: cst.Assign(
+                        # pyre-ignore: Incompatible parameter type [6]
+                        targets=[
+                            cst.BinaryOperation(
+                                left=cst.Name("x"),
+                                operator=cst.Add(),
+                                right=cst.Integer("1"),
+                            ),
+                        ],
+                        value=cst.Name("y"),
+                    )
+                ),
+                "expected_re": "Expected an instance of .*statement.AssignTarget.*",
+            },
+        )
+    )
+    def test_invalid_types(self, **kwargs: Any) -> None:
+        self.assert_invalid_types(**kwargs)
+
 
 class AnnAssignTest(CSTNodeTest):
     @data_provider(
@@ -284,6 +307,31 @@ def test_valid(self, **kwargs: Any) -> None:
     def test_invalid(self, **kwargs: Any) -> None:
         self.assert_invalid(**kwargs)
 
+    @data_provider(
+        (
+            {
+                "get_node": (
+                    lambda: cst.AnnAssign(
+                        # pyre-ignore: Incompatible parameter type [6]
+                        target=cst.BinaryOperation(
+                            left=cst.Name("x"),
+                            operator=cst.Add(),
+                            right=cst.Integer("1"),
+                        ),
+                        annotation=cst.Annotation(cst.Name("int")),
+                        equal=cst.AssignEqual(),
+                        value=cst.Name("y"),
+                    )
+                ),
+                "expected_re": (
+                    "Expected an instance of .*BaseAssignTargetExpression.*"
+                ),
+            },
+        )
+    )
+    def test_invalid_types(self, **kwargs: Any) -> None:
+        self.assert_invalid_types(**kwargs)
+
 
 class AugAssignTest(CSTNodeTest):
     @data_provider(
@@ -362,3 +410,27 @@ class AugAssignTest(CSTNodeTest):
     )
     def test_valid(self, **kwargs: Any) -> None:
         self.validate_node(**kwargs)
+
+    @data_provider(
+        (
+            {
+                "get_node": (
+                    lambda: cst.AugAssign(
+                        # pyre-ignore: Incompatible parameter type [6]
+                        target=cst.BinaryOperation(
+                            left=cst.Name("x"),
+                            operator=cst.Add(),
+                            right=cst.Integer("1"),
+                        ),
+                        operator=cst.Add(),
+                        value=cst.Name("y"),
+                    )
+                ),
+                "expected_re": (
+                    "Expected an instance of .*BaseAssignTargetExpression.*"
+                ),
+            },
+        )
+    )
+    def test_invalid_types(self, **kwargs: Any) -> None:
+        self.assert_invalid_types(**kwargs)
diff --git a/libcst/matchers/__init__.py b/libcst/matchers/__init__.py
@@ -223,10 +223,10 @@ class And(BaseBooleanOp, BaseMatcherNode):
     ] = DoNotCare()
 
 
-BaseExpressionMatchType = Union[
-    "BaseExpression",
+BaseAssignTargetExpressionMatchType = Union[
+    "BaseAssignTargetExpression",
     MetadataMatchType,
-    MatchIfTrue[Callable[[cst.BaseExpression], bool]],
+    MatchIfTrue[Callable[[cst.BaseAssignTargetExpression], bool]],
 ]
 AnnotationMatchType = Union[
     "Annotation", MetadataMatchType, MatchIfTrue[Callable[[cst.Annotation], bool]]
@@ -248,10 +248,10 @@ class And(BaseBooleanOp, BaseMatcherNode):
 @dataclass(frozen=True, eq=False, unsafe_hash=False)
 class AnnAssign(BaseSmallStatement, BaseMatcherNode):
     target: Union[
-        BaseExpressionMatchType,
+        BaseAssignTargetExpressionMatchType,
         DoNotCareSentinel,
-        OneOf[BaseExpressionMatchType],
-        AllOf[BaseExpressionMatchType],
+        OneOf[BaseAssignTargetExpressionMatchType],
+        AllOf[BaseAssignTargetExpressionMatchType],
     ] = DoNotCare()
     annotation: Union[
         AnnotationMatchType,
@@ -285,6 +285,13 @@ class AnnAssign(BaseSmallStatement, BaseMatcherNode):
     ] = DoNotCare()
 
 
+BaseExpressionMatchType = Union[
+    "BaseExpression",
+    MetadataMatchType,
+    MatchIfTrue[Callable[[cst.BaseExpression], bool]],
+]
+
+
 @dataclass(frozen=True, eq=False, unsafe_hash=False)
 class Annotation(BaseMatcherNode):
     annotation: Union[
@@ -597,13 +604,6 @@ class AssignEqual(BaseMatcherNode):
     ] = DoNotCare()
 
 
-BaseAssignTargetExpressionMatchType = Union[
-    "BaseAssignTargetExpression",
-    MetadataMatchType,
-    MatchIfTrue[Callable[[cst.BaseAssignTargetExpression], bool]],
-]
-
-
 @dataclass(frozen=True, eq=False, unsafe_hash=False)
 class AssignTarget(BaseMatcherNode):
     target: Union[
@@ -852,10 +852,10 @@ class Attribute(
 @dataclass(frozen=True, eq=False, unsafe_hash=False)
 class AugAssign(BaseSmallStatement, BaseMatcherNode):
     target: Union[
-        BaseExpressionMatchType,
+        BaseAssignTargetExpressionMatchType,
         DoNotCareSentinel,
-        OneOf[BaseExpressionMatchType],
-        AllOf[BaseExpressionMatchType],
+        OneOf[BaseAssignTargetExpressionMatchType],
+        AllOf[BaseAssignTargetExpressionMatchType],
     ] = DoNotCare()
     operator: Union[
         BaseAugOpMatchType,