Fix space validation for AsName and Await (#641)

* Fix space validation for AsName and Await * Update libcst/_nodes/tests/test_import.py Co-authored-by: Steven Troxler <[email protected]>
Instagram · Feb 10, 2022 · 3af6820 · 3af6820
1 parent 1aa40f7
commit 3af6820
Show file tree

Hide file tree

Showing 5 changed files with 93 additions and 13 deletions.
diff --git a/libcst/_nodes/expression.py b/libcst/_nodes/expression.py
@@ -2394,7 +2394,12 @@ def _validate(self) -> None:
         # Validate any super-class stuff, whatever it may be.
         super(Await, self)._validate()
         # Make sure we don't run identifiers together.
-        if self.whitespace_after_await.empty:
+        if (
+            self.whitespace_after_await.empty
+            and not self.expression._safe_to_use_with_word_operator(
+                ExpressionPosition.RIGHT
+            )
+        ):
             raise CSTValidationError("Must have at least one space after await")
 
     def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Await":

diff --git a/libcst/_nodes/statement.py b/libcst/_nodes/statement.py
@@ -750,8 +750,6 @@ def _validate(self) -> None:
             raise CSTValidationError(
                 "There must be at least one space between 'as' and name."
             )
-        if self.whitespace_before_as.empty:
-            raise CSTValidationError("There must be at least one space before 'as'.")
 
     def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "AsName":
         return AsName(
@@ -815,6 +813,16 @@ def _validate(self) -> None:
                 raise CSTValidationError(
                     "Must have at least one space after except when ExceptHandler has a type."
                 )
+        name = self.name
+        if (
+            type_ is not None
+            and name is not None
+            and name.whitespace_before_as.empty
+            and not type_._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
+        ):
+            raise CSTValidationError(
+                "Must have at least one space before as keyword in an except handler."
+            )
 
     def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "ExceptHandler":
         return ExceptHandler(
@@ -1139,10 +1147,15 @@ class ImportAlias(CSTNode):
 
     def _validate(self) -> None:
         asname = self.asname
-        if asname is not None and not isinstance(asname.name, Name):
-            raise CSTValidationError(
-                "Must use a Name node for AsName name inside ImportAlias."
-            )
+        if asname is not None:
+            if not isinstance(asname.name, Name):
+                raise CSTValidationError(
+                    "Must use a Name node for AsName name inside ImportAlias."
+                )
+            if asname.whitespace_before_as.empty:
+                raise CSTValidationError(
+                    "Must have at least one space before as keyword in an ImportAlias."
+                )
         try:
             self.evaluated_name
         except Exception as e:
@@ -1986,6 +1999,15 @@ class WithItem(CSTNode):
     #: other items inside a with block must contain a comma to separate them.
     comma: Union[Comma, MaybeSentinel] = MaybeSentinel.DEFAULT
 
+    def _validate(self) -> None:
+        asname = self.asname
+        if (
+            asname is not None
+            and asname.whitespace_before_as.empty
+            and not self.item._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
+        ):
+            raise CSTValidationError("Must have at least one space before as keyword.")
+
     def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "WithItem":
         return WithItem(
             item=visit_required(self, "item", self.item, visitor),

diff --git a/libcst/_nodes/tests/test_await.py b/libcst/_nodes/tests/test_await.py
@@ -46,6 +46,14 @@ class AwaitTest(CSTNodeTest):
                 ),
                 "expected_position": CodeRange((1, 2), (1, 13)),
             },
+            # Whitespace after await
+            {
+                "node": cst.Await(
+                    cst.Name("foo", lpar=[cst.LeftParen()], rpar=[cst.RightParen()]),
+                    whitespace_after_await=cst.SimpleWhitespace(""),
+                ),
+                "code": "await(foo)",
+            },
         )
     )
     def test_valid_py37(self, **kwargs: Any) -> None:

diff --git a/libcst/_nodes/tests/test_import.py b/libcst/_nodes/tests/test_import.py
@@ -195,6 +195,20 @@ def test_valid(self, **kwargs: Any) -> None:
                 ),
                 "expected_re": "at least one space",
             },
+            {
+                "get_node": lambda: cst.Import(
+                    names=(
+                        cst.ImportAlias(
+                            cst.Name("foo"),
+                            asname=cst.AsName(
+                                cst.Name("bar"),
+                                whitespace_before_as=cst.SimpleWhitespace(""),
+                            ),
+                        ),
+                    ),
+                ),
+                "expected_re": "at least one space",
+            },
             {
                 "get_node": lambda: cst.Import(
                     names=[
@@ -564,6 +578,25 @@ def test_valid(self, **kwargs: Any) -> None:
                 ),
                 "expected_re": "one space after import",
             },
+            {
+                "get_node": lambda: cst.ImportFrom(
+                    module=cst.Name("foo"),
+                    names=(
+                        cst.ImportAlias(
+                            cst.Name("bar"),
+                            asname=cst.AsName(
+                                cst.Name(
+                                    "baz",
+                                    lpar=(cst.LeftParen(),),
+                                    rpar=(cst.RightParen(),),
+                                ),
+                                whitespace_before_as=cst.SimpleWhitespace(""),
+                            ),
+                        ),
+                    ),
+                ),
+                "expected_re": "one space before as keyword",
+            },
         )
     )
     def test_invalid(self, **kwargs: Any) -> None:

diff --git a/libcst/_nodes/tests/test_try.py b/libcst/_nodes/tests/test_try.py
@@ -329,6 +329,24 @@ class TryTest(CSTNodeTest):
                 "code": "try: pass\nexcept(IOError, ImportError): pass\n",
                 "parser": parse_statement,
             },
+            # No space before as
+            {
+                "node": cst.Try(
+                    cst.SimpleStatementSuite((cst.Pass(),)),
+                    handlers=[
+                        cst.ExceptHandler(
+                            cst.SimpleStatementSuite((cst.Pass(),)),
+                            whitespace_after_except=cst.SimpleWhitespace(" "),
+                            type=cst.Call(cst.Name("foo")),
+                            name=cst.AsName(
+                                whitespace_before_as=cst.SimpleWhitespace(""),
+                                name=cst.Name("bar"),
+                            ),
+                        )
+                    ],
+                ),
+                "code": "try: pass\nexcept foo()as bar: pass\n",
+            },
         )
     )
     def test_valid(self, **kwargs: Any) -> None:
@@ -346,12 +364,6 @@ def test_valid(self, **kwargs: Any) -> None:
                 ),
                 "expected_re": "between 'as'",
             },
-            {
-                "get_node": lambda: cst.AsName(
-                    cst.Name("bla"), whitespace_before_as=cst.SimpleWhitespace("")
-                ),
-                "expected_re": "before 'as'",
-            },
             {
                 "get_node": lambda: cst.ExceptHandler(
                     cst.SimpleStatementSuite((cst.Pass(),)),