mirror of
https://github.com/Instagram/LibCST.git
synced 2025-12-23 10:35:53 +00:00
Define DictComp node
This defines the node and adds tests for it, but doesn't implement the parser for it. That will come in a later PR.
This commit is contained in:
Benjamin Woodruff
Benjamin Woodruff
parent
27b8cab777
commit
cf2dfa5ee6
4 changed files with 230 additions and 6 deletions
|
|
@ -45,7 +45,7 @@ Expressions
|
|||
.. autoclass:: libcst.CompIf
|
||||
.. autoclass:: libcst.ConcatenatedString
|
||||
.. autoclass:: libcst.Dict
|
||||
.. .. autoclass:: libcst.DictComp
|
||||
.. autoclass:: libcst.DictComp
|
||||
.. autoclass:: libcst.DictElement
|
||||
.. autoclass:: libcst.Element
|
||||
.. autoclass:: libcst.Ellipses
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ from libcst._nodes._expression import (
|
|||
CompIf,
|
||||
ConcatenatedString,
|
||||
Dict,
|
||||
DictComp,
|
||||
DictElement,
|
||||
Element,
|
||||
Ellipses,
|
||||
|
|
@ -228,6 +229,7 @@ __all__ = [
|
|||
"CompIf",
|
||||
"ConcatenatedString",
|
||||
"Dict",
|
||||
"DictComp",
|
||||
"DictElement",
|
||||
"Element",
|
||||
"Ellipses",
|
||||
|
|
|
|||
|
|
@ -2832,14 +2832,15 @@ class BaseSimpleComp(BaseComp, ABC):
|
|||
`BaseSimpleComp`, because it uses `key` and `value`.
|
||||
"""
|
||||
|
||||
# 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.
|
||||
#: 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.
|
||||
# pyre-fixme[13]: Attribute `elt` is never initialized.
|
||||
elt: BaseAssignTargetExpression
|
||||
|
||||
# The `for ... in ... if ...` clause that lexically comes after `elt`. This may be a
|
||||
# nested structure for nested comprehensions. See `ComprehensionFor` for details.
|
||||
#: The ``for ... in ... if ...`` clause that lexically comes after ``elt``. This may
|
||||
#: be a nested structure for nested comprehensions. See :class:`.CompFor` for
|
||||
#: details.
|
||||
# pyre-fixme[13]: Attribute `for_in` is never initialized.
|
||||
for_in: CompFor
|
||||
|
||||
|
|
@ -2941,3 +2942,72 @@ class SetComp(BaseSet, BaseSimpleComp):
|
|||
with self._parenthesize(state), self._braceize(state):
|
||||
self.elt._codegen(state)
|
||||
self.for_in._codegen(state)
|
||||
|
||||
|
||||
@add_slots
|
||||
@dataclass(frozen=True)
|
||||
class DictComp(BaseDict, BaseComp):
|
||||
"""
|
||||
A dictionary comprehension. ``key`` and ``value`` represent the dictionary entry
|
||||
evaluated for each item.
|
||||
|
||||
All ``for ... in ...`` and ``if ...`` clauses are stored as a recursive
|
||||
:class:`CompFor` data structure inside ``for_in``.
|
||||
"""
|
||||
|
||||
key: BaseAssignTargetExpression
|
||||
value: BaseAssignTargetExpression
|
||||
|
||||
#: The ``for ... in ... if ...`` clause that lexically comes after ``key`` and
|
||||
#: ``value``. This may be a nested structure for nested comprehensions. See
|
||||
#: :class:`.CompFor` for details.
|
||||
for_in: CompFor
|
||||
|
||||
lbrace: LeftCurlyBrace = LeftCurlyBrace()
|
||||
rbrace: RightCurlyBrace = RightCurlyBrace()
|
||||
lpar: Sequence[LeftParen] = ()
|
||||
rpar: Sequence[RightParen] = ()
|
||||
|
||||
#: Whitespace after the key, but before the colon in ``key : value``.
|
||||
whitespace_before_colon: BaseParenthesizableWhitespace = SimpleWhitespace("")
|
||||
#: Whitespace after the colon, but before the value in ``key : value``.
|
||||
whitespace_after_colon: BaseParenthesizableWhitespace = SimpleWhitespace(" ")
|
||||
|
||||
def _validate(self) -> None:
|
||||
super(DictComp, self)._validate()
|
||||
|
||||
for_in = self.for_in
|
||||
if (
|
||||
for_in.whitespace_before.empty
|
||||
and not self.value._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
|
||||
):
|
||||
keyword = "async" if for_in.asynchronous else "for"
|
||||
raise CSTValidationError(
|
||||
f"Must have at least one space before '{keyword}' keyword."
|
||||
)
|
||||
|
||||
def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "DictComp":
|
||||
return DictComp(
|
||||
lpar=visit_sequence("lpar", self.lpar, visitor),
|
||||
lbrace=visit_required("lbrace", self.lbrace, visitor),
|
||||
key=visit_required("key", self.key, visitor),
|
||||
whitespace_before_colon=visit_required(
|
||||
"whitespace_before_colon", self.whitespace_before_colon, visitor
|
||||
),
|
||||
whitespace_after_colon=visit_required(
|
||||
"whitespace_after_colon", self.whitespace_after_colon, visitor
|
||||
),
|
||||
value=visit_required("value", self.value, visitor),
|
||||
for_in=visit_required("for_in", self.for_in, visitor),
|
||||
rbrace=visit_required("rbrace", self.rbrace, visitor),
|
||||
rpar=visit_sequence("rpar", self.rpar, visitor),
|
||||
)
|
||||
|
||||
def _codegen_impl(self, state: CodegenState) -> None:
|
||||
with self._parenthesize(state), self._braceize(state):
|
||||
self.key._codegen(state)
|
||||
self.whitespace_before_colon._codegen(state)
|
||||
state.add_token(":")
|
||||
self.whitespace_after_colon._codegen(state)
|
||||
self.value._codegen(state)
|
||||
self.for_in._codegen(state)
|
||||
|
|
|
|||
152
libcst/_nodes/tests/test_dict_comp.py
Normal file
152
libcst/_nodes/tests/test_dict_comp.py
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# pyre-strict
|
||||
from typing import Any
|
||||
|
||||
import libcst as cst
|
||||
from libcst import CodeRange
|
||||
from libcst._nodes.tests.base import CSTNodeTest
|
||||
from libcst.testing.utils import data_provider
|
||||
|
||||
|
||||
class DictCompTest(CSTNodeTest):
|
||||
@data_provider(
|
||||
[
|
||||
# simple DictComp
|
||||
{
|
||||
"node": cst.DictComp(
|
||||
cst.Name("k"),
|
||||
cst.Name("v"),
|
||||
cst.CompFor(target=cst.Name("a"), iter=cst.Name("b")),
|
||||
),
|
||||
"code": "{k: v for a in b}",
|
||||
"expected_position": CodeRange.create((1, 0), (1, 17)),
|
||||
},
|
||||
# custom whitespace around colon
|
||||
{
|
||||
"node": cst.DictComp(
|
||||
cst.Name("k"),
|
||||
cst.Name("v"),
|
||||
cst.CompFor(target=cst.Name("a"), iter=cst.Name("b")),
|
||||
whitespace_before_colon=cst.SimpleWhitespace("\t"),
|
||||
whitespace_after_colon=cst.SimpleWhitespace("\t\t"),
|
||||
),
|
||||
"code": "{k\t:\t\tv for a in b}",
|
||||
"expected_position": CodeRange.create((1, 0), (1, 19)),
|
||||
},
|
||||
# custom whitespace inside braces
|
||||
{
|
||||
"node": cst.DictComp(
|
||||
cst.Name("k"),
|
||||
cst.Name("v"),
|
||||
cst.CompFor(target=cst.Name("a"), iter=cst.Name("b")),
|
||||
lbrace=cst.LeftCurlyBrace(
|
||||
whitespace_after=cst.SimpleWhitespace("\t")
|
||||
),
|
||||
rbrace=cst.RightCurlyBrace(
|
||||
whitespace_before=cst.SimpleWhitespace("\t\t")
|
||||
),
|
||||
),
|
||||
"code": "{\tk: v for a in b\t\t}",
|
||||
"expected_position": CodeRange.create((1, 0), (1, 20)),
|
||||
},
|
||||
# parenthesis
|
||||
{
|
||||
"node": cst.DictComp(
|
||||
cst.Name("k"),
|
||||
cst.Name("v"),
|
||||
cst.CompFor(target=cst.Name("a"), iter=cst.Name("b")),
|
||||
lpar=[cst.LeftParen()],
|
||||
rpar=[cst.RightParen()],
|
||||
),
|
||||
"code": "({k: v for a in b})",
|
||||
"expected_position": CodeRange.create((1, 1), (1, 18)),
|
||||
},
|
||||
# missing spaces around DictComp is always okay
|
||||
{
|
||||
"node": cst.DictComp(
|
||||
cst.Name("a"),
|
||||
cst.Name("b"),
|
||||
cst.CompFor(
|
||||
target=cst.Name("c"),
|
||||
iter=cst.DictComp(
|
||||
cst.Name("d"),
|
||||
cst.Name("e"),
|
||||
cst.CompFor(target=cst.Name("f"), iter=cst.Name("g")),
|
||||
),
|
||||
ifs=[
|
||||
cst.CompIf(
|
||||
cst.Name("h"),
|
||||
whitespace_before=cst.SimpleWhitespace(""),
|
||||
)
|
||||
],
|
||||
whitespace_after_in=cst.SimpleWhitespace(""),
|
||||
),
|
||||
),
|
||||
"code": "{a: b for c in{d: e for f in g}if h}",
|
||||
"expected_position": CodeRange.create((1, 0), (1, 36)),
|
||||
},
|
||||
# no whitespace before `for` clause
|
||||
{
|
||||
"node": cst.DictComp(
|
||||
cst.Name("k"),
|
||||
cst.Name("v", lpar=[cst.LeftParen()], rpar=[cst.RightParen()]),
|
||||
cst.CompFor(
|
||||
target=cst.Name("a"),
|
||||
iter=cst.Name("b"),
|
||||
whitespace_before=cst.SimpleWhitespace(""),
|
||||
),
|
||||
),
|
||||
"code": "{k: (v)for a in b}",
|
||||
"expected_position": CodeRange.create((1, 0), (1, 18)),
|
||||
},
|
||||
]
|
||||
)
|
||||
def test_valid(self, **kwargs: Any) -> None:
|
||||
self.validate_node(**kwargs)
|
||||
|
||||
@data_provider(
|
||||
[
|
||||
# unbalanced DictComp
|
||||
{
|
||||
"get_node": lambda: cst.DictComp(
|
||||
cst.Name("k"),
|
||||
cst.Name("v"),
|
||||
cst.CompFor(target=cst.Name("a"), iter=cst.Name("b")),
|
||||
lpar=[cst.LeftParen()],
|
||||
),
|
||||
"expected_re": "left paren without right paren",
|
||||
},
|
||||
# invalid whitespace before for/async
|
||||
{
|
||||
"get_node": lambda: cst.DictComp(
|
||||
cst.Name("k"),
|
||||
cst.Name("v"),
|
||||
cst.CompFor(
|
||||
target=cst.Name("a"),
|
||||
iter=cst.Name("b"),
|
||||
whitespace_before=cst.SimpleWhitespace(""),
|
||||
),
|
||||
),
|
||||
"expected_re": "Must have at least one space before 'for' keyword.",
|
||||
},
|
||||
{
|
||||
"get_node": lambda: cst.DictComp(
|
||||
cst.Name("k"),
|
||||
cst.Name("v"),
|
||||
cst.CompFor(
|
||||
target=cst.Name("a"),
|
||||
iter=cst.Name("b"),
|
||||
asynchronous=cst.Asynchronous(),
|
||||
whitespace_before=cst.SimpleWhitespace(""),
|
||||
),
|
||||
),
|
||||
"expected_re": "Must have at least one space before 'async' keyword.",
|
||||
},
|
||||
]
|
||||
)
|
||||
def test_invalid(self, **kwargs: Any) -> None:
|
||||
self.assert_invalid(**kwargs)
|
||||
Loading…
Add table
Add a link
Reference in a new issue