Record syntactic position of more expressions

Update codegen and tests to cover list, list comprehension and generator
expressions. Also cleans up some extraneous position tracking code in
codegen.

libcst/nodes/_expression.py

@@ -206,7 +206,8 @@ class Asynchronous(CSTNode):
         )
 
     def _codegen_impl(self, state: CodegenState) -> None:
-        state.add_token("async")
+        with state.record_syntactic_position(self):
+            state.add_token("async")
         self.whitespace_after._codegen(state)
 
 
@@ -598,21 +599,20 @@ class FormattedStringExpression(BaseFormattedStringContent):
         )
 
     def _codegen_impl(self, state: CodegenState) -> None:
-        with state.record_syntactic_position(self):
-            state.add_token("{")
-            self.whitespace_before_expression._codegen(state)
-            self.expression._codegen(state)
-            self.whitespace_after_expression._codegen(state)
-            conversion = self.conversion
-            if conversion is not None:
-                state.add_token("!")
-                state.add_token(conversion)
-            format_spec = self.format_spec
-            if format_spec is not None:
-                state.add_token(":")
-                for spec in format_spec:
-                    spec._codegen(state)
-            state.add_token("}")
+        state.add_token("{")
+        self.whitespace_before_expression._codegen(state)
+        self.expression._codegen(state)
+        self.whitespace_after_expression._codegen(state)
+        conversion = self.conversion
+        if conversion is not None:
+            state.add_token("!")
+            state.add_token(conversion)
+        format_spec = self.format_spec
+        if format_spec is not None:
+            state.add_token(":")
+            for spec in format_spec:
+                spec._codegen(state)
+        state.add_token("}")
 
 
 @add_slots
@@ -1033,8 +1033,7 @@ class Index(CSTNode):
         return Index(value=visit_required("value", self.value, visitor))
 
     def _codegen_impl(self, state: CodegenState) -> None:
-        with state.record_syntactic_position(self):
-            self.value._codegen(state)
+        self.value._codegen(state)
 
 
 @add_slots
@@ -1071,22 +1070,21 @@ class Slice(CSTNode):
         )
 
     def _codegen_impl(self, state: CodegenState) -> None:
-        with state.record_syntactic_position(self):
-            lower = self.lower
-            if lower is not None:
-                lower._codegen(state)
-            self.first_colon._codegen(state)
-            upper = self.upper
-            if upper is not None:
-                upper._codegen(state)
-            second_colon = self.second_colon
-            if second_colon is MaybeSentinel.DEFAULT and self.step is not None:
-                state.add_token(":")
-            elif isinstance(second_colon, Colon):
-                second_colon._codegen(state)
-            step = self.step
-            if step is not None:
-                step._codegen(state)
+        lower = self.lower
+        if lower is not None:
+            lower._codegen(state)
+        self.first_colon._codegen(state)
+        upper = self.upper
+        if upper is not None:
+            upper._codegen(state)
+        second_colon = self.second_colon
+        if second_colon is MaybeSentinel.DEFAULT and self.step is not None:
+            state.add_token(":")
+        elif isinstance(second_colon, Colon):
+            second_colon._codegen(state)
+        step = self.step
+        if step is not None:
+            step._codegen(state)
 
 
 @add_slots
@@ -1476,62 +1474,58 @@ class Parameters(CSTNode):
         )
 
     def _codegen_impl(self, state: CodegenState) -> None:
-        # TODO: remove this when fallback to syntactic whitespace becomes available
-        with state.record_syntactic_position(self):
-            # Compute the star existence first so we can ask about whether
-            # each element is the last in the list or not.
-            star_arg = self.star_arg
-            if isinstance(star_arg, MaybeSentinel):
-                starincluded = len(self.kwonly_params) > 0
-            elif isinstance(star_arg, (Param, ParamStar)):
-                starincluded = True
-            else:
-                starincluded = False
-            # Render out the params first, computing necessary trailing commas.
-            lastparam = len(self.params) - 1
-            more_values = (
-                len(self.default_params) > 0
-                or starincluded
-                or len(self.kwonly_params) > 0
-                or self.star_kwarg is not None
+        # Compute the star existence first so we can ask about whether
+        # each element is the last in the list or not.
+        star_arg = self.star_arg
+        if isinstance(star_arg, MaybeSentinel):
+            starincluded = len(self.kwonly_params) > 0
+        elif isinstance(star_arg, (Param, ParamStar)):
+            starincluded = True
+        else:
+            starincluded = False
+        # Render out the params first, computing necessary trailing commas.
+        lastparam = len(self.params) - 1
+        more_values = (
+            len(self.default_params) > 0
+            or starincluded
+            or len(self.kwonly_params) > 0
+            or self.star_kwarg is not None
+        )
+        for i, param in enumerate(self.params):
+            param._codegen(
+                state, default_star="", default_comma=(i < lastparam or more_values)
             )
-            for i, param in enumerate(self.params):
-                param._codegen(
-                    state, default_star="", default_comma=(i < lastparam or more_values)
-                )
-            # Render out the default_params next, computing necessary trailing commas.
-            lastparam = len(self.default_params) - 1
-            more_values = (
-                starincluded
-                or len(self.kwonly_params) > 0
-                or self.star_kwarg is not None
+        # Render out the default_params next, computing necessary trailing commas.
+        lastparam = len(self.default_params) - 1
+        more_values = (
+            starincluded or len(self.kwonly_params) > 0 or self.star_kwarg is not None
+        )
+        for i, param in enumerate(self.default_params):
+            param._codegen(
+                state, default_star="", default_comma=(i < lastparam or more_values)
             )
-            for i, param in enumerate(self.default_params):
-                param._codegen(
-                    state, default_star="", default_comma=(i < lastparam or more_values)
-                )
-            # Render out optional star sentinel if its explicitly included or
-            # if we are inferring it from kwonly_params. Otherwise, render out the
-            # optional star_arg.
-            if isinstance(star_arg, MaybeSentinel):
-                if starincluded:
-                    state.add_token("*, ")
-            elif isinstance(star_arg, Param):
-                more_values = len(self.kwonly_params) > 0 or self.star_kwarg is not None
-                star_arg._codegen(state, default_star="*", default_comma=more_values)
-            elif isinstance(star_arg, ParamStar):
-                star_arg._codegen(state)
-            # Render out the kwonly_args next, computing necessary trailing commas.
-            lastparam = len(self.kwonly_params) - 1
-            more_values = self.star_kwarg is not None
-            for i, param in enumerate(self.kwonly_params):
-                param._codegen(
-                    state, default_star="", default_comma=(i < lastparam or more_values)
-                )
-            # Finally, render out any optional star_kwarg
-            star_kwarg = self.star_kwarg
-            if star_kwarg is not None:
-                star_kwarg._codegen(state, default_star="**", default_comma=False)
+        # Render out optional star sentinel if its explicitly included or
+        # if we are inferring it from kwonly_params. Otherwise, render out the
+        # optional star_arg.
+        if isinstance(star_arg, MaybeSentinel):
+            if starincluded:
+                state.add_token("*, ")
+        elif isinstance(star_arg, Param):
+            more_values = len(self.kwonly_params) > 0 or self.star_kwarg is not None
+            star_arg._codegen(state, default_star="*", default_comma=more_values)
+        elif isinstance(star_arg, ParamStar):
+            star_arg._codegen(state)
+        # Render out the kwonly_args next, computing necessary trailing commas.
+        lastparam = len(self.kwonly_params) - 1
+        more_values = self.star_kwarg is not None
+        for i, param in enumerate(self.kwonly_params):
+            param._codegen(
+                state, default_star="", default_comma=(i < lastparam or more_values)
+            )
+        # Finally, render out any optional star_kwarg
+        star_kwarg = self.star_kwarg
+        if star_kwarg is not None:
+            star_kwarg._codegen(state, default_star="**", default_comma=False)
 
 
 @add_slots
@@ -1839,15 +1833,14 @@ class Call(_BaseExpressionWithArgs):
 
     def _codegen_impl(self, state: CodegenState) -> None:
         with self._parenthesize(state):
-            with state.record_syntactic_position(self):
-                self.func._codegen(state)
-                self.whitespace_after_func._codegen(state)
-                state.add_token("(")
-                self.whitespace_before_args._codegen(state)
-                lastarg = len(self.args) - 1
-                for i, arg in enumerate(self.args):
-                    arg._codegen(state, default_comma=(i != lastarg))
-                state.add_token(")")
+            self.func._codegen(state)
+            self.whitespace_after_func._codegen(state)
+            state.add_token("(")
+            self.whitespace_before_args._codegen(state)
+            lastarg = len(self.args) - 1
+            for i, arg in enumerate(self.args):
+                arg._codegen(state, default_comma=(i != lastarg))
+            state.add_token(")")
 
 
 @add_slots
@@ -2144,6 +2137,7 @@ class Element(BaseElement):
     ) -> None:
         with state.record_syntactic_position(self):
             self.value._codegen(state)
+
         comma = self.comma
         if comma is MaybeSentinel.DEFAULT and default_comma:
             if default_comma_whitespace:

libcst/nodes/_internal.py

@@ -118,7 +118,13 @@ class CodegenState:
             node._metadata[self.provider] = position
 
     @contextmanager
-    def record_syntactic_position(self, node: _CSTNodeT) -> Iterator[None]:
+    def record_syntactic_position(
+        self,
+        node: _CSTNodeT,
+        *,
+        start_node: Optional[_CSTNodeT] = None,
+        end_node: Optional[_CSTNodeT] = None,
+    ) -> Iterator[None]:
         yield
 
 
@@ -133,12 +139,27 @@ class SyntacticCodegenState(CodegenState):
         self.provider = SyntacticPositionProvider
 
     @contextmanager
-    def record_syntactic_position(self, node: _CSTNodeT) -> Iterator[None]:
+    def record_syntactic_position(
+        self,
+        node: _CSTNodeT,
+        *,
+        start_node: Optional[_CSTNodeT] = None,
+        end_node: Optional[_CSTNodeT] = None,
+    ) -> Iterator[None]:
         start = CodePosition(self.line, self.column)
         try:
             yield
         finally:
             end = CodePosition(self.line, self.column)
+
+            # Override with positions hoisted from child nodes if provided
+            start = (
+                start_node._metadata[self.provider].start
+                if start_node is not None
+                else start
+            )
+            end = end_node._metadata[self.provider].end if end_node is not None else end
+
             node._metadata[self.provider] = CodeRange(start, end)
 
 

libcst/nodes/tests/test_list.py

@@ -7,6 +7,7 @@
 from typing import Any, Callable
 
 import libcst.nodes as cst
+from libcst.nodes._internal import CodeRange
 from libcst.nodes.tests.base import CSTNodeTest
 from libcst.parser import parse_expression, parse_statement
 from libcst.testing.utils import data_provider
@@ -39,6 +40,7 @@ class ListTest(CSTNodeTest):
                 ),
                 "code": "[\tsingle_element    ]",
                 "parser": parse_expression,
+                "expected_position": CodeRange.create((1, 0), (1, 21)),
             },
             # two-element list, sentinel comma value
             {
@@ -57,6 +59,7 @@ class ListTest(CSTNodeTest):
                 ),
                 "code": "([one])",
                 "parser": None,
+                "expected_position": CodeRange.create((1, 1), (1, 6)),
             },
             # starred element
             {
@@ -68,6 +71,7 @@ class ListTest(CSTNodeTest):
                 ),
                 "code": "[*one, *two]",
                 "parser": None,
+                "expected_position": CodeRange.create((1, 0), (1, 12)),
             },
             # missing spaces around list, always okay
             {

libcst/nodes/tests/test_simple_comp.py

@@ -7,6 +7,7 @@
 from typing import Any, Callable
 
 import libcst.nodes as cst
+from libcst.nodes._internal import CodeRange
 from libcst.nodes.tests.base import CSTNodeTest
 from libcst.parser import parse_expression
 from libcst.testing.utils import data_provider
@@ -22,6 +23,7 @@ class SimpleCompTest(CSTNodeTest):
                 ),
                 "code": "(a for b in c)",
                 "parser": parse_expression,
+                "expected_position": CodeRange.create((1, 1), (1, 13)),
             },
             # simple ListComp
             {
@@ -30,6 +32,7 @@ class SimpleCompTest(CSTNodeTest):
                 ),
                 "code": "[a for b in c]",
                 "parser": parse_expression,
+                "expected_position": CodeRange.create((1, 0), (1, 14)),
             },
             # simple SetComp
             {
@@ -86,6 +89,7 @@ class SimpleCompTest(CSTNodeTest):
                 ),
                 "code": "(a for b in c if d if e if f)",
                 "parser": parse_expression,
+                "expected_position": CodeRange.create((1, 1), (1, 28)),
             },
             # nested/inner for-in clause
             {
@@ -144,6 +148,7 @@ class SimpleCompTest(CSTNodeTest):
                 ),
                 "code": "(\fa  for   b    in     c\tif\t\td\f\f)",
                 "parser": parse_expression,
+                "expected_position": CodeRange.create((1, 2), (1, 30)),
             },
             # custom whitespace around ListComp's brackets
             {
@@ -163,6 +168,7 @@ class SimpleCompTest(CSTNodeTest):
                 ),
                 "code": "(\f[\ta for b in c\t\t]\f\f)",
                 "parser": parse_expression,
+                "expected_position": CodeRange.create((1, 2), (1, 19)),
             },
             # custom whitespace around SetComp's braces
             {
@@ -225,6 +231,7 @@ class SimpleCompTest(CSTNodeTest):
                 ),
                 "code": "((a)for(b)in(c)if(d)for(e)in(f))",
                 "parser": parse_expression,
+                "expected_position": CodeRange.create((1, 1), (1, 31)),
             },
             # no whitespace before/after GeneratorExp is valid
             {