Fix whitespace, fstring, walrus related parse errors (#939, #938, #937, #936, #935, #934, #933, #932, #931)

* Allow walrus in slices

See python/cpython#23317

Raised in #930.

* Fix parsing of nested f-string specifiers

For an expression like `f"{one:{two:}{three}}"`, `three` is not in an f-string spec, and should be tokenized accordingly.

This PR fixes the `format_spec_count` bookkeeping in the tokenizer, so it properly decrements it when a closing `}` is encountered but only if the `}` closes a format_spec.

Reported in #930.

* Fix tokenizing `0else`

This is an obscure one.

`_ if 0else _` failed to parse with some very weird errors. It turns out that the tokenizer tries to parse `0else` as a single number, but when it encounters `l` it realizes it can't be a single number and it backtracks.

Unfortunately the backtracking logic was broken, and it failed to correctly backtrack one of the offsets used for whitespace parsing (the byte offset since the start of the line). This caused whitespace nodes to refer to incorrect parts of the input text, eventually resulting in the above behavior.

This PR fixes the bookkeeping when the tokenizer backtracks.

Reported in #930.

* Allow no whitespace between lambda keyword and params in certain cases

Python accepts code where `lambda` follows a `*`, so this PR relaxes validation rules for Lambdas.

Raised in #930.

* Allow any expression in comprehensions' evaluated expression


This PR relaxes the accepted types for the `elt` field in `ListComp`, `SetComp`, and `GenExp`, as well as the `key` and `value` fields in `DictComp`.

Fixes #500.

* Allow no space around an ifexp in certain cases

For example in `_ if _ else""if _ else _`.

Raised in #930. Also fixes #854.

* Allow no spaces after `as` in a contextmanager in certain cases

Like in `with foo()as():pass`

Raised in #930.

* Allow no spaces around walrus in certain cases

Like in `[_:=''for _ in _]`

Raised in #930.

* Allow no whitespace after lambda body in certain cases

Like in `[lambda:()for _ in _]`

Reported in #930.

Author: zsol

libcst/_nodes/expression.py

@@ -1983,6 +1983,25 @@ def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Parameters":
             star_kwarg=visit_optional(self, "star_kwarg", self.star_kwarg, visitor),
         )
 
+    def _safe_to_join_with_lambda(self) -> bool:
+        """
+        Determine if Parameters need a space after the `lambda` keyword. Returns True
+        iff it's safe to omit the space between `lambda` and these Parameters.
+
+        See also `BaseExpression._safe_to_use_with_word_operator`.
+
+        For example: `lambda*_: pass`
+        """
+        if len(self.posonly_params) != 0:
+            return False
+
+        # posonly_ind can't appear if above condition is false
+
+        if len(self.params) > 0 and self.params[0].star not in {"*", "**"}:
+            return False
+
+        return True
+
     def _codegen_impl(self, state: CodegenState) -> None:  # noqa: C901
         # Compute the star existence first so we can ask about whether
         # each element is the last in the list or not.
@@ -2088,6 +2107,13 @@ class Lambda(BaseExpression):
         BaseParenthesizableWhitespace, MaybeSentinel
     ] = MaybeSentinel.DEFAULT
 
+    def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
+        if position == ExpressionPosition.LEFT:
+            return len(self.rpar) > 0 or self.body._safe_to_use_with_word_operator(
+                position
+            )
+        return super()._safe_to_use_with_word_operator(position)
+
     def _validate(self) -> None:
         # Validate parents
         super(Lambda, self)._validate()
@@ -2115,6 +2141,7 @@ def _validate(self) -> None:
             if (
                 isinstance(whitespace_after_lambda, BaseParenthesizableWhitespace)
                 and whitespace_after_lambda.empty
+                and not self.params._safe_to_join_with_lambda()
             ):
                 raise CSTValidationError(
                     "Must have at least one space after lambda when specifying params"
@@ -2492,6 +2519,12 @@ class IfExp(BaseExpression):
     #: Whitespace after the ``else`` keyword, but before the ``orelse`` expression.
     whitespace_after_else: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
 
+    def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
+        if position == ExpressionPosition.RIGHT:
+            return self.body._safe_to_use_with_word_operator(position)
+        else:
+            return self.orelse._safe_to_use_with_word_operator(position)
+
     def _validate(self) -> None:
         # Paren validation and such
         super(IfExp, self)._validate()
@@ -3495,7 +3528,7 @@ class BaseSimpleComp(BaseComp, ABC):
     #: The expression evaluated during each iteration of the comprehension. This
     #: lexically comes before the ``for_in`` clause, but it is semantically the
     #: inner-most element, evaluated inside the ``for_in`` clause.
-    elt: BaseAssignTargetExpression
+    elt: BaseExpression
 
     #: The ``for ... in ... if ...`` clause that lexically comes after ``elt``. This may
     #: be a nested structure for nested comprehensions. See :class:`CompFor` for
@@ -3528,7 +3561,7 @@ class GeneratorExp(BaseSimpleComp):
     """
 
     #: The expression evaluated and yielded during each iteration of the generator.
-    elt: BaseAssignTargetExpression
+    elt: BaseExpression
 
     #: The ``for ... in ... if ...`` clause that comes after ``elt``. This may be a
     #: nested structure for nested comprehensions. See :class:`CompFor` for details.
@@ -3579,7 +3612,7 @@ class ListComp(BaseList, BaseSimpleComp):
     """
 
     #: The expression evaluated and stored during each iteration of the comprehension.
-    elt: BaseAssignTargetExpression
+    elt: BaseExpression
 
     #: The ``for ... in ... if ...`` clause that comes after ``elt``. This may be a
     #: nested structure for nested comprehensions. See :class:`CompFor` for details.
@@ -3621,7 +3654,7 @@ class SetComp(BaseSet, BaseSimpleComp):
     """
 
     #: The expression evaluated and stored during each iteration of the comprehension.
-    elt: BaseAssignTargetExpression
+    elt: BaseExpression
 
     #: The ``for ... in ... if ...`` clause that comes after ``elt``. This may be a
     #: nested structure for nested comprehensions. See :class:`CompFor` for details.
@@ -3663,10 +3696,10 @@ class DictComp(BaseDict, BaseComp):
     """
 
     #: The key inserted into the dictionary during each iteration of the comprehension.
-    key: BaseAssignTargetExpression
+    key: BaseExpression
     #: The value associated with the ``key`` inserted into the dictionary during each
     #: iteration of the comprehension.
-    value: BaseAssignTargetExpression
+    value: BaseExpression
 
     #: The ``for ... in ... if ...`` clause that lexically comes after ``key`` and
     #: ``value``. This may be a nested structure for nested comprehensions. See
@@ -3770,6 +3803,15 @@ def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "NamedExpr":
             rpar=visit_sequence(self, "rpar", self.rpar, visitor),
         )
 
+    def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
+        if position == ExpressionPosition.LEFT:
+            return len(self.rpar) > 0 or self.value._safe_to_use_with_word_operator(
+                position
+            )
+        return len(self.lpar) > 0 or self.target._safe_to_use_with_word_operator(
+            position
+        )
+
     def _codegen_impl(self, state: CodegenState) -> None:
         with self._parenthesize(state):
             self.target._codegen(state)

libcst/_nodes/statement.py

@@ -745,7 +745,10 @@ class AsName(CSTNode):
     whitespace_after_as: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
 
     def _validate(self) -> None:
-        if self.whitespace_after_as.empty:
+        if (
+            self.whitespace_after_as.empty
+            and not self.name._safe_to_use_with_word_operator(ExpressionPosition.RIGHT)
+        ):
             raise CSTValidationError(
                 "There must be at least one space between 'as' and name."
             )

libcst/_nodes/tests/test_dict_comp.py

@@ -26,6 +26,17 @@ class DictCompTest(CSTNodeTest):
                 "parser": parse_expression,
                 "expected_position": CodeRange((1, 0), (1, 17)),
             },
+            # non-trivial keys & values in DictComp
+            {
+                "node": cst.DictComp(
+                    cst.BinaryOperation(cst.Name("k1"), cst.Add(), cst.Name("k2")),
+                    cst.BinaryOperation(cst.Name("v1"), cst.Add(), cst.Name("v2")),
+                    cst.CompFor(target=cst.Name("a"), iter=cst.Name("b")),
+                ),
+                "code": "{k1 + k2: v1 + v2 for a in b}",
+                "parser": parse_expression,
+                "expected_position": CodeRange((1, 0), (1, 29)),
+            },
             # custom whitespace around colon
             {
                 "node": cst.DictComp(

libcst/_nodes/tests/test_ifexp.py

@@ -52,6 +52,41 @@ class IfExpTest(CSTNodeTest):
                 "(foo)if(bar)else(baz)",
                 CodeRange((1, 0), (1, 21)),
             ),
+            (
+                cst.IfExp(
+                    body=cst.Name("foo"),
+                    whitespace_before_if=cst.SimpleWhitespace(" "),
+                    whitespace_after_if=cst.SimpleWhitespace(" "),
+                    test=cst.Name("bar"),
+                    whitespace_before_else=cst.SimpleWhitespace(" "),
+                    whitespace_after_else=cst.SimpleWhitespace(""),
+                    orelse=cst.IfExp(
+                        body=cst.SimpleString("''"),
+                        whitespace_before_if=cst.SimpleWhitespace(""),
+                        test=cst.Name("bar"),
+                        orelse=cst.Name("baz"),
+                    ),
+                ),
+                "foo if bar else''if bar else baz",
+                CodeRange((1, 0), (1, 32)),
+            ),
+            (
+                cst.GeneratorExp(
+                    elt=cst.IfExp(
+                        body=cst.Name("foo"),
+                        test=cst.Name("bar"),
+                        orelse=cst.SimpleString("''"),
+                        whitespace_after_else=cst.SimpleWhitespace(""),
+                    ),
+                    for_in=cst.CompFor(
+                        target=cst.Name("_"),
+                        iter=cst.Name("_"),
+                        whitespace_before=cst.SimpleWhitespace(""),
+                    ),
+                ),
+                "(foo if bar else''for _ in _)",
+                CodeRange((1, 1), (1, 28)),
+            ),
             # Make sure that spacing works
             (
                 cst.IfExp(

libcst/_nodes/tests/test_lambda.py

@@ -303,30 +303,6 @@ def test_valid(
                 ),
                 "at least one space after lambda",
             ),
-            (
-                lambda: cst.Lambda(
-                    cst.Parameters(star_arg=cst.Param(cst.Name("arg"))),
-                    cst.Integer("5"),
-                    whitespace_after_lambda=cst.SimpleWhitespace(""),
-                ),
-                "at least one space after lambda",
-            ),
-            (
-                lambda: cst.Lambda(
-                    cst.Parameters(kwonly_params=(cst.Param(cst.Name("arg")),)),
-                    cst.Integer("5"),
-                    whitespace_after_lambda=cst.SimpleWhitespace(""),
-                ),
-                "at least one space after lambda",
-            ),
-            (
-                lambda: cst.Lambda(
-                    cst.Parameters(star_kwarg=cst.Param(cst.Name("arg"))),
-                    cst.Integer("5"),
-                    whitespace_after_lambda=cst.SimpleWhitespace(""),
-                ),
-                "at least one space after lambda",
-            ),
             (
                 lambda: cst.Lambda(
                     cst.Parameters(
@@ -944,6 +920,53 @@ class LambdaParserTest(CSTNodeTest):
                 ),
                 "( lambda  : 5 )",
             ),
+            # No space between lambda and params
+            (
+                cst.Lambda(
+                    cst.Parameters(star_arg=cst.Param(cst.Name("args"), star="*")),
+                    cst.Integer("5"),
+                    whitespace_after_lambda=cst.SimpleWhitespace(""),
+                ),
+                "lambda*args: 5",
+            ),
+            (
+                cst.Lambda(
+                    cst.Parameters(star_kwarg=cst.Param(cst.Name("kwargs"), star="**")),
+                    cst.Integer("5"),
+                    whitespace_after_lambda=cst.SimpleWhitespace(""),
+                ),
+                "lambda**kwargs: 5",
+            ),
+            (
+                cst.Lambda(
+                    cst.Parameters(
+                        star_arg=cst.ParamStar(
+                            comma=cst.Comma(
+                                cst.SimpleWhitespace(""), cst.SimpleWhitespace("")
+                            )
+                        ),
+                        kwonly_params=[cst.Param(cst.Name("args"), star="")],
+                    ),
+                    cst.Integer("5"),
+                    whitespace_after_lambda=cst.SimpleWhitespace(""),
+                ),
+                "lambda*,args: 5",
+            ),
+            (
+                cst.ListComp(
+                    elt=cst.Lambda(
+                        params=cst.Parameters(),
+                        body=cst.Tuple(()),
+                        colon=cst.Colon(),
+                    ),
+                    for_in=cst.CompFor(
+                        target=cst.Name("_"),
+                        iter=cst.Name("_"),
+                        whitespace_before=cst.SimpleWhitespace(""),
+                    ),
+                ),
+                "[lambda:()for _ in _]",
+            ),
         )
     )
     def test_valid(

libcst/_nodes/tests/test_namedexpr.py

@@ -166,6 +166,22 @@ class NamedExprTest(CSTNodeTest):
                 "parser": _parse_expression_force_38,
                 "expected_position": None,
             },
+            {
+                "node": cst.ListComp(
+                    elt=cst.NamedExpr(
+                        cst.Name("_"),
+                        cst.SimpleString("''"),
+                        whitespace_after_walrus=cst.SimpleWhitespace(""),
+                        whitespace_before_walrus=cst.SimpleWhitespace(""),
+                    ),
+                    for_in=cst.CompFor(
+                        target=cst.Name("_"),
+                        iter=cst.Name("_"),
+                        whitespace_before=cst.SimpleWhitespace(""),
+                    ),
+                ),
+                "code": "[_:=''for _ in _]",
+            },
         )
     )
     def test_valid(self, **kwargs: Any) -> None:

libcst/_nodes/tests/test_simple_comp.py

@@ -41,6 +41,33 @@ class SimpleCompTest(CSTNodeTest):
                 "code": "{a for b in c}",
                 "parser": parse_expression,
             },
+            # non-trivial elt in GeneratorExp
+            {
+                "node": cst.GeneratorExp(
+                    cst.BinaryOperation(cst.Name("a1"), cst.Add(), cst.Name("a2")),
+                    cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
+                ),
+                "code": "(a1 + a2 for b in c)",
+                "parser": parse_expression,
+            },
+            # non-trivial elt in ListComp
+            {
+                "node": cst.ListComp(
+                    cst.BinaryOperation(cst.Name("a1"), cst.Add(), cst.Name("a2")),
+                    cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
+                ),
+                "code": "[a1 + a2 for b in c]",
+                "parser": parse_expression,
+            },
+            # non-trivial elt in SetComp
+            {
+                "node": cst.SetComp(
+                    cst.BinaryOperation(cst.Name("a1"), cst.Add(), cst.Name("a2")),
+                    cst.CompFor(target=cst.Name("b"), iter=cst.Name("c")),
+                ),
+                "code": "{a1 + a2 for b in c}",
+                "parser": parse_expression,
+            },
             # async GeneratorExp
             {
                 "node": cst.GeneratorExp(

libcst/_nodes/tests/test_with.py

@@ -102,6 +102,23 @@ class WithTest(CSTNodeTest):
                 "code": "with context_mgr() as ctx: pass\n",
                 "parser": parse_statement,
             },
+            {
+                "node": cst.With(
+                    (
+                        cst.WithItem(
+                            cst.Call(cst.Name("context_mgr")),
+                            cst.AsName(
+                                cst.Tuple(()),
+                                whitespace_after_as=cst.SimpleWhitespace(""),
+                                whitespace_before_as=cst.SimpleWhitespace(""),
+                            ),
+                        ),
+                    ),
+                    cst.SimpleStatementSuite((cst.Pass(),)),
+                ),
+                "code": "with context_mgr()as(): pass\n",
+                "parser": parse_statement,
+            },
             # indentation
             {
                 "node": DummyIndentedBlock(