1# Copyright (c) Meta Platforms, Inc. and affiliates.
2#
3# This source code is licensed under the MIT license found in the
4# LICENSE file in the root directory of this source tree.
5
6
7import re
8from abc import ABC, abstractmethod
9from ast import literal_eval
10from contextlib import contextmanager
11from dataclasses import dataclass, field
12from enum import auto, Enum
13from tokenize import (
14 Floatnumber as FLOATNUMBER_RE,
15 Imagnumber as IMAGNUMBER_RE,
16 Intnumber as INTNUMBER_RE,
17)
18from typing import Callable, Generator, Literal, Optional, Sequence, Union
19
20from libcst import CSTLogicError
21
22from libcst._add_slots import add_slots
23from libcst._maybe_sentinel import MaybeSentinel
24from libcst._nodes.base import CSTCodegenError, CSTNode, CSTValidationError
25from libcst._nodes.internal import (
26 CodegenState,
27 visit_optional,
28 visit_required,
29 visit_sentinel,
30 visit_sequence,
31)
32from libcst._nodes.op import (
33 AssignEqual,
34 BaseBinaryOp,
35 BaseBooleanOp,
36 BaseCompOp,
37 BaseUnaryOp,
38 Colon,
39 Comma,
40 Dot,
41 In,
42 Is,
43 IsNot,
44 Not,
45 NotIn,
46)
47from libcst._nodes.whitespace import BaseParenthesizableWhitespace, SimpleWhitespace
48from libcst._visitors import CSTVisitorT
49
50
51@add_slots
52@dataclass(frozen=True)
53class LeftSquareBracket(CSTNode):
54 """
55 Used by various nodes to denote a subscript or list section. This doesn't own
56 the whitespace to the left of it since this is owned by the parent node.
57 """
58
59 #: Any space that appears directly after this left square bracket.
60 whitespace_after: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
61
62 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "LeftSquareBracket":
63 return LeftSquareBracket(
64 whitespace_after=visit_required(
65 self, "whitespace_after", self.whitespace_after, visitor
66 )
67 )
68
69 def _codegen_impl(self, state: CodegenState) -> None:
70 state.add_token("[")
71 self.whitespace_after._codegen(state)
72
73
74@add_slots
75@dataclass(frozen=True)
76class RightSquareBracket(CSTNode):
77 """
78 Used by various nodes to denote a subscript or list section. This doesn't own
79 the whitespace to the right of it since this is owned by the parent node.
80 """
81
82 #: Any space that appears directly before this right square bracket.
83 whitespace_before: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
84
85 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "RightSquareBracket":
86 return RightSquareBracket(
87 whitespace_before=visit_required(
88 self, "whitespace_before", self.whitespace_before, visitor
89 )
90 )
91
92 def _codegen_impl(self, state: CodegenState) -> None:
93 self.whitespace_before._codegen(state)
94 state.add_token("]")
95
96
97@add_slots
98@dataclass(frozen=True)
99class LeftCurlyBrace(CSTNode):
100 """
101 Used by various nodes to denote a dict or set. This doesn't own the whitespace to
102 the left of it since this is owned by the parent node.
103 """
104
105 #: Any space that appears directly after this left curly brace.
106 whitespace_after: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
107
108 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "LeftCurlyBrace":
109 return LeftCurlyBrace(
110 whitespace_after=visit_required(
111 self, "whitespace_after", self.whitespace_after, visitor
112 )
113 )
114
115 def _codegen_impl(self, state: CodegenState) -> None:
116 state.add_token("{")
117 self.whitespace_after._codegen(state)
118
119
120@add_slots
121@dataclass(frozen=True)
122class RightCurlyBrace(CSTNode):
123 """
124 Used by various nodes to denote a dict or set. This doesn't own the whitespace to
125 the right of it since this is owned by the parent node.
126 """
127
128 #: Any space that appears directly before this right curly brace.
129 whitespace_before: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
130
131 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "RightCurlyBrace":
132 return RightCurlyBrace(
133 whitespace_before=visit_required(
134 self, "whitespace_before", self.whitespace_before, visitor
135 )
136 )
137
138 def _codegen_impl(self, state: CodegenState) -> None:
139 self.whitespace_before._codegen(state)
140 state.add_token("}")
141
142
143@add_slots
144@dataclass(frozen=True)
145class LeftParen(CSTNode):
146 """
147 Used by various nodes to denote a parenthesized section. This doesn't own
148 the whitespace to the left of it since this is owned by the parent node.
149 """
150
151 #: Any space that appears directly after this left parenthesis.
152 whitespace_after: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
153
154 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "LeftParen":
155 return LeftParen(
156 whitespace_after=visit_required(
157 self, "whitespace_after", self.whitespace_after, visitor
158 )
159 )
160
161 def _codegen_impl(self, state: CodegenState) -> None:
162 state.add_token("(")
163 self.whitespace_after._codegen(state)
164
165
166@add_slots
167@dataclass(frozen=True)
168class RightParen(CSTNode):
169 """
170 Used by various nodes to denote a parenthesized section. This doesn't own
171 the whitespace to the right of it since this is owned by the parent node.
172 """
173
174 #: Any space that appears directly after this left parenthesis.
175 whitespace_before: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
176
177 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "RightParen":
178 return RightParen(
179 whitespace_before=visit_required(
180 self, "whitespace_before", self.whitespace_before, visitor
181 )
182 )
183
184 def _codegen_impl(self, state: CodegenState) -> None:
185 self.whitespace_before._codegen(state)
186 state.add_token(")")
187
188
189@add_slots
190@dataclass(frozen=True)
191class Asynchronous(CSTNode):
192 """
193 Used by asynchronous function definitions, as well as ``async for`` and
194 ``async with``.
195 """
196
197 #: Any space that appears directly after this async keyword.
198 whitespace_after: SimpleWhitespace = SimpleWhitespace.field(" ")
199
200 def _validate(self) -> None:
201 if len(self.whitespace_after.value) < 1:
202 raise CSTValidationError("Must have at least one space after Asynchronous.")
203
204 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Asynchronous":
205 return Asynchronous(
206 whitespace_after=visit_required(
207 self, "whitespace_after", self.whitespace_after, visitor
208 )
209 )
210
211 def _codegen_impl(self, state: CodegenState) -> None:
212 with state.record_syntactic_position(self):
213 state.add_token("async")
214 self.whitespace_after._codegen(state)
215
216
217class _BaseParenthesizedNode(CSTNode, ABC):
218 """
219 We don't want to have another level of indirection for parenthesis in
220 our tree, since that makes us more of a CST than an AST. So, all the
221 expressions or atoms that can be wrapped in parenthesis will subclass
222 this to get that functionality.
223 """
224
225 __slots__ = ()
226
227 lpar: Sequence[LeftParen] = ()
228 # Sequence of parenthesis for precedence dictation.
229 rpar: Sequence[RightParen] = ()
230
231 def _validate(self) -> None:
232 if self.lpar and not self.rpar:
233 raise CSTValidationError("Cannot have left paren without right paren.")
234 if not self.lpar and self.rpar:
235 raise CSTValidationError("Cannot have right paren without left paren.")
236 if len(self.lpar) != len(self.rpar):
237 raise CSTValidationError("Cannot have unbalanced parens.")
238
239 @contextmanager
240 def _parenthesize(self, state: CodegenState) -> Generator[None, None, None]:
241 for lpar in self.lpar:
242 lpar._codegen(state)
243 with state.record_syntactic_position(self):
244 yield
245 for rpar in self.rpar:
246 rpar._codegen(state)
247
248
249class ExpressionPosition(Enum):
250 LEFT = auto()
251 RIGHT = auto()
252
253
254class BaseExpression(_BaseParenthesizedNode, ABC):
255 """
256 An base class for all expressions. :class:`BaseExpression` contains no fields.
257 """
258
259 __slots__ = ()
260
261 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
262 """
263 Returns true if this expression is safe to be use with a word operator
264 such as "not" without space between the operator an ourselves. Examples
265 where this is true are "not(True)", "(1)in[1,2,3]", etc. This base
266 function handles parenthesized nodes, but certain nodes such as tuples,
267 dictionaries and lists will override this to signifiy that they're always
268 safe.
269 """
270
271 return len(self.lpar) > 0 and len(self.rpar) > 0
272
273 def _check_left_right_word_concatenation_safety(
274 self,
275 position: ExpressionPosition,
276 left: "BaseExpression",
277 right: "BaseExpression",
278 ) -> bool:
279 if position == ExpressionPosition.RIGHT:
280 return left._safe_to_use_with_word_operator(ExpressionPosition.RIGHT)
281 if position == ExpressionPosition.LEFT:
282 return right._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
283 return False
284
285
286class BaseAssignTargetExpression(BaseExpression, ABC):
287 """
288 An expression that's valid on the left side of an assignment. That assignment may
289 be part an :class:`Assign` node, or it may be part of a number of other control
290 structures that perform an assignment, such as a :class:`For` loop.
291
292 Python's grammar defines all expression as valid in this position, but the AST
293 compiler further restricts the allowed types, which is what this type attempts to
294 express.
295
296 This is similar to a :class:`BaseDelTargetExpression`, but it also includes
297 :class:`StarredElement` as a valid node.
298
299 The set of valid nodes are defined as part of `CPython's AST context computation
300 <https://github.com/python/cpython/blob/v3.8.0a4/Python/ast.c#L1120>`_.
301 """
302
303 __slots__ = ()
304
305
306class BaseDelTargetExpression(BaseExpression, ABC):
307 """
308 An expression that's valid on the right side of a :class:`Del` statement.
309
310 Python's grammar defines all expression as valid in this position, but the AST
311 compiler further restricts the allowed types, which is what this type attempts to
312 express.
313
314 This is similar to a :class:`BaseAssignTargetExpression`, but it excludes
315 :class:`StarredElement`.
316
317 The set of valid nodes are defined as part of `CPython's AST context computation
318 <https://github.com/python/cpython/blob/v3.8.0a4/Python/ast.c#L1120>`_ and as part
319 of `CPython's bytecode compiler
320 <https://github.com/python/cpython/blob/v3.8.0a4/Python/compile.c#L4854>`_.
321 """
322
323 __slots__ = ()
324
325
326@add_slots
327@dataclass(frozen=True)
328class Name(BaseAssignTargetExpression, BaseDelTargetExpression):
329 """
330 A simple variable name. Names are typically used in the context of a variable
331 access, an assignment, or a deletion.
332
333 Dotted variable names (``a.b.c``) are represented with :class:`Attribute` nodes,
334 and subscripted variable names (``a[b]``) are represented with :class:`Subscript`
335 nodes.
336 """
337
338 #: The variable's name (or "identifier") as a string.
339 value: str
340
341 lpar: Sequence[LeftParen] = ()
342 #: Sequence of parenthesis for precedence dictation.
343 rpar: Sequence[RightParen] = ()
344
345 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Name":
346 return Name(
347 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
348 value=self.value,
349 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
350 )
351
352 def _validate(self) -> None:
353 super(Name, self)._validate()
354 if len(self.value) == 0:
355 raise CSTValidationError("Cannot have empty name identifier.")
356 if not self.value.isidentifier():
357 raise CSTValidationError(f"Name {self.value!r} is not a valid identifier.")
358
359 def _codegen_impl(self, state: CodegenState) -> None:
360 with self._parenthesize(state):
361 state.add_token(self.value)
362
363
364@add_slots
365@dataclass(frozen=True)
366class Ellipsis(BaseExpression):
367 """
368 An ellipsis ``...``. When used as an expression, it evaluates to the
369 `Ellipsis constant`_. Ellipsis are often used as placeholders in code or in
370 conjunction with :class:`SubscriptElement`.
371
372 .. _Ellipsis constant: https://docs.python.org/3/library/constants.html#Ellipsis
373 """
374
375 lpar: Sequence[LeftParen] = ()
376
377 #: Sequence of parenthesis for precedence dictation.
378 rpar: Sequence[RightParen] = ()
379
380 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Ellipsis":
381 return Ellipsis(
382 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
383 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
384 )
385
386 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
387 return True
388
389 def _codegen_impl(self, state: CodegenState) -> None:
390 with self._parenthesize(state):
391 state.add_token("...")
392
393
394class BaseNumber(BaseExpression, ABC):
395 """
396 A type such as :class:`Integer`, :class:`Float`, or :class:`Imaginary` that can be
397 used anywhere that you need to explicitly take any number type.
398 """
399
400 __slots__ = ()
401
402 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
403 """
404 Numbers are funny. The expression "5in [1,2,3,4,5]" is a valid expression
405 which evaluates to "True". So, encapsulate that here by allowing zero spacing
406 with the left hand side of an expression with a comparison operator.
407 """
408 if position == ExpressionPosition.LEFT:
409 return True
410 return super(BaseNumber, self)._safe_to_use_with_word_operator(position)
411
412
413@add_slots
414@dataclass(frozen=True)
415class Integer(BaseNumber):
416 #: A string representation of the integer, such as ``"100000"`` or ``100_000``.
417 #:
418 #: To convert this string representation to an ``int``, use the calculated
419 #: property :attr:`~Integer.evaluated_value`.
420 value: str
421
422 lpar: Sequence[LeftParen] = ()
423 #: Sequence of parenthesis for precedence dictation.
424 rpar: Sequence[RightParen] = ()
425
426 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Integer":
427 return Integer(
428 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
429 value=self.value,
430 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
431 )
432
433 def _validate(self) -> None:
434 super(Integer, self)._validate()
435 if not re.fullmatch(INTNUMBER_RE, self.value):
436 raise CSTValidationError("Number is not a valid integer.")
437
438 def _codegen_impl(self, state: CodegenState) -> None:
439 with self._parenthesize(state):
440 state.add_token(self.value)
441
442 @property
443 def evaluated_value(self) -> int:
444 """
445 Return an :func:`ast.literal_eval` evaluated int of :py:attr:`value`.
446 """
447 return literal_eval(self.value)
448
449
450@add_slots
451@dataclass(frozen=True)
452class Float(BaseNumber):
453 #: A string representation of the floating point number, such as ``"0.05"``,
454 #: ``".050"``, or ``"5e-2"``.
455 #:
456 #: To convert this string representation to an ``float``, use the calculated
457 #: property :attr:`~Float.evaluated_value`.
458 value: str
459
460 lpar: Sequence[LeftParen] = ()
461 #: Sequence of parenthesis for precedence dictation.
462 rpar: Sequence[RightParen] = ()
463
464 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Float":
465 return Float(
466 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
467 value=self.value,
468 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
469 )
470
471 def _validate(self) -> None:
472 super(Float, self)._validate()
473 if not re.fullmatch(FLOATNUMBER_RE, self.value):
474 raise CSTValidationError("Number is not a valid float.")
475
476 def _codegen_impl(self, state: CodegenState) -> None:
477 with self._parenthesize(state):
478 state.add_token(self.value)
479
480 @property
481 def evaluated_value(self) -> float:
482 """
483 Return an :func:`ast.literal_eval` evaluated float of :py:attr:`value`.
484 """
485 return literal_eval(self.value)
486
487
488@add_slots
489@dataclass(frozen=True)
490class Imaginary(BaseNumber):
491 #: A string representation of the imaginary (complex) number, such as ``"2j"``.
492 #:
493 #: To convert this string representation to an ``complex``, use the calculated
494 #: property :attr:`~Imaginary.evaluated_value`.
495 value: str
496
497 lpar: Sequence[LeftParen] = ()
498 #: Sequence of parenthesis for precedence dictation.
499 rpar: Sequence[RightParen] = ()
500
501 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Imaginary":
502 return Imaginary(
503 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
504 value=self.value,
505 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
506 )
507
508 def _validate(self) -> None:
509 super(Imaginary, self)._validate()
510 if not re.fullmatch(IMAGNUMBER_RE, self.value):
511 raise CSTValidationError("Number is not a valid imaginary.")
512
513 def _codegen_impl(self, state: CodegenState) -> None:
514 with self._parenthesize(state):
515 state.add_token(self.value)
516
517 @property
518 def evaluated_value(self) -> complex:
519 """
520 Return an :func:`ast.literal_eval` evaluated complex of :py:attr:`value`.
521 """
522 return literal_eval(self.value)
523
524
525class BaseString(BaseExpression, ABC):
526 """
527 A type that can be used anywhere that you need to take any string. This includes
528 :class:`SimpleString`, :class:`ConcatenatedString`, and :class:`FormattedString`.
529 """
530
531 __slots__ = ()
532
533
534StringQuoteLiteral = Literal['"', "'", '"""', "'''"]
535
536
537class _BasePrefixedString(BaseString, ABC):
538 __slots__ = ()
539
540 @property
541 def prefix(self) -> str:
542 """
543 Returns the string's prefix, if any exists.
544
545 See `String and Bytes literals
546 <https://docs.python.org/3.7/reference/lexical_analysis.html#string-and-bytes-literals>`_
547 for more information.
548 """
549 ...
550
551 @property
552 def quote(self) -> StringQuoteLiteral:
553 """
554 Returns the quotation used to denote the string. Can be either ``'``,
555 ``"``, ``'''`` or ``\"\"\"``.
556 """
557 ...
558
559 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
560 """
561 ``"a"in"abc`` is okay, but if you add a prefix, (e.g. ``b"a"inb"abc"``), the string
562 is no longer valid on the RHS of the word operator, because it's not clear where
563 the keyword ends and the prefix begins, unless it's parenthesized.
564 """
565 if position == ExpressionPosition.LEFT:
566 return True
567 elif self.prefix == "": # and position == ExpressionPosition.RIGHT
568 return True
569 else:
570 return super(_BasePrefixedString, self)._safe_to_use_with_word_operator(
571 position
572 )
573
574
575@add_slots
576@dataclass(frozen=True)
577class SimpleString(_BasePrefixedString):
578 """
579 Any sort of literal string expression that is not a :class:`FormattedString`
580 (f-string), including triple-quoted multi-line strings.
581 """
582
583 #: The texual representation of the string, including quotes, prefix characters, and
584 #: any escape characters present in the original source code , such as
585 #: ``r"my string\n"``. To remove the quotes and interpret any escape characters,
586 #: use the calculated property :attr:`~SimpleString.evaluated_value`.
587 value: str
588
589 lpar: Sequence[LeftParen] = ()
590 #: Sequence of parenthesis for precidence dictation.
591 rpar: Sequence[RightParen] = ()
592
593 def _validate(self) -> None:
594 super(SimpleString, self)._validate()
595
596 # Validate any prefix
597 prefix = self.prefix
598 if prefix not in ("", "r", "u", "b", "br", "rb"):
599 raise CSTValidationError("Invalid string prefix.")
600 prefixlen = len(prefix)
601 # Validate wrapping quotes
602 if len(self.value) < (prefixlen + 2):
603 raise CSTValidationError("String must have enclosing quotes.")
604 if (
605 self.value[prefixlen] not in ['"', "'"]
606 or self.value[prefixlen] != self.value[-1]
607 ):
608 raise CSTValidationError("String must have matching enclosing quotes.")
609 # Check validity of triple-quoted strings
610 if len(self.value) >= (prefixlen + 6):
611 if self.value[prefixlen] == self.value[prefixlen + 1]:
612 # We know this isn't an empty string, so there needs to be a third
613 # identical enclosing token.
614 if (
615 self.value[prefixlen] != self.value[prefixlen + 2]
616 or self.value[prefixlen] != self.value[-2]
617 or self.value[prefixlen] != self.value[-3]
618 ):
619 raise CSTValidationError(
620 "String must have matching enclosing quotes."
621 )
622 # We should check the contents as well, but this is pretty complicated,
623 # partially due to triple-quoted strings.
624
625 @property
626 def prefix(self) -> str:
627 """
628 Returns the string's prefix, if any exists. The prefix can be ``r``,
629 ``u``, ``b``, ``br`` or ``rb``.
630 """
631
632 prefix: str = ""
633 for c in self.value:
634 if c in ['"', "'"]:
635 break
636 prefix += c
637 return prefix.lower()
638
639 @property
640 def quote(self) -> StringQuoteLiteral:
641 """
642 Returns the quotation used to denote the string. Can be either ``'``,
643 ``"``, ``'''`` or ``\"\"\"``.
644 """
645
646 quote: str = ""
647 for char in self.value[len(self.prefix) :]:
648 if char not in {"'", '"'}:
649 break
650 if quote and char != quote[0]:
651 # This is no longer the same string quote
652 break
653 quote += char
654
655 if len(quote) == 2:
656 # Let's assume this is an empty string.
657 quote = quote[:1]
658 elif 3 < len(quote) <= 6:
659 # Let's assume this can be one of the following:
660 # >>> """"foo"""
661 # '"foo'
662 # >>> """""bar"""
663 # '""bar'
664 # >>> """"""
665 # ''
666 quote = quote[:3]
667
668 if len(quote) not in {1, 3}:
669 # We shouldn't get here due to construction validation logic,
670 # but handle the case anyway.
671 raise CSTLogicError(f"Invalid string {self.value}")
672
673 # pyre-ignore We know via the above validation that we will only
674 # ever return one of the four string literals.
675 return quote
676
677 @property
678 def raw_value(self) -> str:
679 """
680 Returns the raw value of the string as it appears in source, without
681 the beginning or end quotes and without the prefix. This is often
682 useful when constructing transforms which need to manipulate strings
683 in source code.
684 """
685
686 prefix_len = len(self.prefix)
687 quote_len = len(self.quote)
688 return self.value[(prefix_len + quote_len) : (-quote_len)]
689
690 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "SimpleString":
691 return SimpleString(
692 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
693 value=self.value,
694 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
695 )
696
697 def _codegen_impl(self, state: CodegenState) -> None:
698 with self._parenthesize(state):
699 state.add_token(self.value)
700
701 @property
702 def evaluated_value(self) -> Union[str, bytes]:
703 """
704 Return an :func:`ast.literal_eval` evaluated str of :py:attr:`value`.
705 """
706 return literal_eval(self.value)
707
708
709class BaseFormattedStringContent(CSTNode, ABC):
710 """
711 The base type for :class:`FormattedStringText` and
712 :class:`FormattedStringExpression`. A :class:`FormattedString` is composed of a
713 sequence of :class:`BaseFormattedStringContent` parts.
714 """
715
716 __slots__ = ()
717
718
719@add_slots
720@dataclass(frozen=True)
721class FormattedStringText(BaseFormattedStringContent):
722 """
723 Part of a :class:`FormattedString` that is not inside curly braces (``{`` or ``}``).
724 For example, in::
725
726 f"ab{cd}ef"
727
728 ``ab`` and ``ef`` are :class:`FormattedStringText` nodes, but ``{cd}`` is a
729 :class:`FormattedStringExpression`.
730 """
731
732 #: The raw string value, including any escape characters present in the source
733 #: code, not including any enclosing quotes.
734 value: str
735
736 def _visit_and_replace_children(
737 self, visitor: CSTVisitorT
738 ) -> "FormattedStringText":
739 return FormattedStringText(value=self.value)
740
741 def _codegen_impl(self, state: CodegenState) -> None:
742 state.add_token(self.value)
743
744
745@add_slots
746@dataclass(frozen=True)
747class FormattedStringExpression(BaseFormattedStringContent):
748 """
749 Part of a :class:`FormattedString` that is inside curly braces (``{`` or ``}``),
750 including the surrounding curly braces. For example, in::
751
752 f"ab{cd}ef"
753
754 ``{cd}`` is a :class:`FormattedStringExpression`, but ``ab`` and ``ef`` are
755 :class:`FormattedStringText` nodes.
756
757 An f-string expression may contain ``conversion`` and ``format_spec`` suffixes that
758 control how the expression is converted to a string. See `Python's language
759 reference
760 <https://docs.python.org/3/reference/lexical_analysis.html#formatted-string-literals>`__
761 for details.
762 """
763
764 #: The expression we will evaluate and render when generating the string.
765 expression: BaseExpression
766
767 #: An optional conversion specifier, such as ``!s``, ``!r`` or ``!a``.
768 conversion: Optional[str] = None
769
770 #: An optional format specifier following the `format specification mini-language
771 #: <https://docs.python.org/3/library/string.html#formatspec>`_.
772 format_spec: Optional[Sequence[BaseFormattedStringContent]] = None
773
774 #: Whitespace after the opening curly brace (``{``), but before the ``expression``.
775 whitespace_before_expression: BaseParenthesizableWhitespace = (
776 SimpleWhitespace.field("")
777 )
778
779 #: Whitespace after the ``expression``, but before the ``conversion``,
780 #: ``format_spec`` and the closing curly brace (``}``). Python does not
781 #: allow whitespace inside or after a ``conversion`` or ``format_spec``.
782 whitespace_after_expression: BaseParenthesizableWhitespace = SimpleWhitespace.field(
783 ""
784 )
785
786 #: Equal sign for formatted string expression uses self-documenting expressions,
787 #: such as ``f"{x=}"``. See the `Python 3.8 release notes
788 #: <https://docs.python.org/3/whatsnew/3.8.html#f-strings-support-for-self-documenting-expressions-and-debugging>`_.
789 equal: Optional[AssignEqual] = None
790
791 def _validate(self) -> None:
792 if self.conversion is not None and self.conversion not in ("s", "r", "a"):
793 raise CSTValidationError("Invalid f-string conversion.")
794
795 def _visit_and_replace_children(
796 self, visitor: CSTVisitorT
797 ) -> "FormattedStringExpression":
798 format_spec = self.format_spec
799 return FormattedStringExpression(
800 whitespace_before_expression=visit_required(
801 self,
802 "whitespace_before_expression",
803 self.whitespace_before_expression,
804 visitor,
805 ),
806 expression=visit_required(self, "expression", self.expression, visitor),
807 equal=visit_optional(self, "equal", self.equal, visitor),
808 whitespace_after_expression=visit_required(
809 self,
810 "whitespace_after_expression",
811 self.whitespace_after_expression,
812 visitor,
813 ),
814 conversion=self.conversion,
815 format_spec=(
816 visit_sequence(self, "format_spec", format_spec, visitor)
817 if format_spec is not None
818 else None
819 ),
820 )
821
822 def _codegen_impl(self, state: CodegenState) -> None:
823 state.add_token("{")
824 self.whitespace_before_expression._codegen(state)
825 self.expression._codegen(state)
826 equal = self.equal
827 if equal is not None:
828 equal._codegen(state)
829 self.whitespace_after_expression._codegen(state)
830 conversion = self.conversion
831 if conversion is not None:
832 state.add_token("!")
833 state.add_token(conversion)
834 format_spec = self.format_spec
835 if format_spec is not None:
836 state.add_token(":")
837 for spec in format_spec:
838 spec._codegen(state)
839 state.add_token("}")
840
841
842@add_slots
843@dataclass(frozen=True)
844class FormattedString(_BasePrefixedString):
845 """
846 An "f-string". These formatted strings are string literals prefixed by the letter
847 "f". An f-string may contain interpolated expressions inside curly braces (``{`` and
848 ``}``).
849
850 F-strings are defined in `PEP 498`_ and documented in `Python's language
851 reference
852 <https://docs.python.org/3/reference/lexical_analysis.html#formatted-string-literals>`__.
853
854 >>> import libcst as cst
855 >>> cst.parse_expression('f"ab{cd}ef"')
856 FormattedString(
857 parts=[
858 FormattedStringText(
859 value='ab',
860 ),
861 FormattedStringExpression(
862 expression=Name(
863 value='cd',
864 lpar=[],
865 rpar=[],
866 ),
867 conversion=None,
868 format_spec=None,
869 whitespace_before_expression=SimpleWhitespace(
870 value='',
871 ),
872 whitespace_after_expression=SimpleWhitespace(
873 value='',
874 ),
875 ),
876 FormattedStringText(
877 value='ef',
878 ),
879 ],
880 start='f"',
881 end='"',
882 lpar=[],
883 rpar=[],
884 )
885
886 .. _PEP 498: https://www.python.org/dev/peps/pep-0498/#specification
887 """
888
889 #: A formatted string is composed as a series of :class:`FormattedStringText` and
890 #: :class:`FormattedStringExpression` parts.
891 parts: Sequence[BaseFormattedStringContent]
892
893 #: The string prefix and the leading quote, such as ``f"``, ``F'``, ``fr"``, or
894 #: ``f"""``.
895 start: str = 'f"'
896
897 #: The trailing quote. This must match the type of quote used in ``start``.
898 end: Literal['"', "'", '"""', "'''"] = '"'
899
900 lpar: Sequence[LeftParen] = ()
901 #: Sequence of parenthesis for precidence dictation.
902 rpar: Sequence[RightParen] = ()
903
904 def _validate(self) -> None:
905 super(FormattedString, self)._validate()
906
907 # Validate any prefix
908 prefix = self.prefix
909 if prefix not in ("f", "fr", "rf"):
910 raise CSTValidationError("Invalid f-string prefix.")
911
912 # Validate wrapping quotes
913 starttoken = self.start[len(prefix) :]
914 if starttoken != self.end:
915 raise CSTValidationError("f-string must have matching enclosing quotes.")
916
917 # Validate valid wrapping quote usage
918 if starttoken not in ('"', "'", '"""', "'''"):
919 raise CSTValidationError("Invalid f-string enclosing quotes.")
920
921 @property
922 def prefix(self) -> str:
923 """
924 Returns the string's prefix, if any exists. The prefix can be ``f``,
925 ``fr``, or ``rf``.
926 """
927
928 prefix = ""
929 for c in self.start:
930 if c in ['"', "'"]:
931 break
932 prefix += c
933 return prefix.lower()
934
935 @property
936 def quote(self) -> StringQuoteLiteral:
937 """
938 Returns the quotation used to denote the string. Can be either ``'``,
939 ``"``, ``'''`` or ``\"\"\"``.
940 """
941
942 return self.end
943
944 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "FormattedString":
945 return FormattedString(
946 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
947 start=self.start,
948 parts=visit_sequence(self, "parts", self.parts, visitor),
949 end=self.end,
950 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
951 )
952
953 def _codegen_impl(self, state: CodegenState) -> None:
954 with self._parenthesize(state):
955 state.add_token(self.start)
956 for part in self.parts:
957 part._codegen(state)
958 state.add_token(self.end)
959
960
961class BaseTemplatedStringContent(CSTNode, ABC):
962 """
963 The base type for :class:`TemplatedStringText` and
964 :class:`TemplatedStringExpression`. A :class:`TemplatedString` is composed of a
965 sequence of :class:`BaseTemplatedStringContent` parts.
966 """
967
968 __slots__ = ()
969
970
971@add_slots
972@dataclass(frozen=True)
973class TemplatedStringText(BaseTemplatedStringContent):
974 """
975 Part of a :class:`TemplatedString` that is not inside curly braces (``{`` or ``}``).
976 For example, in::
977
978 f"ab{cd}ef"
979
980 ``ab`` and ``ef`` are :class:`TemplatedStringText` nodes, but ``{cd}`` is a
981 :class:`TemplatedStringExpression`.
982 """
983
984 #: The raw string value, including any escape characters present in the source
985 #: code, not including any enclosing quotes.
986 value: str
987
988 def _visit_and_replace_children(
989 self, visitor: CSTVisitorT
990 ) -> "TemplatedStringText":
991 return TemplatedStringText(value=self.value)
992
993 def _codegen_impl(self, state: CodegenState) -> None:
994 state.add_token(self.value)
995
996
997@add_slots
998@dataclass(frozen=True)
999class TemplatedStringExpression(BaseTemplatedStringContent):
1000 """
1001 Part of a :class:`TemplatedString` that is inside curly braces (``{`` or ``}``),
1002 including the surrounding curly braces. For example, in::
1003
1004 f"ab{cd}ef"
1005
1006 ``{cd}`` is a :class:`TemplatedStringExpression`, but ``ab`` and ``ef`` are
1007 :class:`TemplatedStringText` nodes.
1008
1009 An t-string expression may contain ``conversion`` and ``format_spec`` suffixes that
1010 control how the expression is converted to a string.
1011 """
1012
1013 #: The expression we will evaluate and render when generating the string.
1014 expression: BaseExpression
1015
1016 #: An optional conversion specifier, such as ``!s``, ``!r`` or ``!a``.
1017 conversion: Optional[str] = None
1018
1019 #: An optional format specifier following the `format specification mini-language
1020 #: <https://docs.python.org/3/library/string.html#formatspec>`_.
1021 format_spec: Optional[Sequence[BaseTemplatedStringContent]] = None
1022
1023 #: Whitespace after the opening curly brace (``{``), but before the ``expression``.
1024 whitespace_before_expression: BaseParenthesizableWhitespace = (
1025 SimpleWhitespace.field("")
1026 )
1027
1028 #: Whitespace after the ``expression``, but before the ``conversion``,
1029 #: ``format_spec`` and the closing curly brace (``}``). Python does not
1030 #: allow whitespace inside or after a ``conversion`` or ``format_spec``.
1031 whitespace_after_expression: BaseParenthesizableWhitespace = SimpleWhitespace.field(
1032 ""
1033 )
1034
1035 #: Equal sign for Templated string expression uses self-documenting expressions,
1036 #: such as ``f"{x=}"``. See the `Python 3.8 release notes
1037 #: <https://docs.python.org/3/whatsnew/3.8.html#f-strings-support-for-self-documenting-expressions-and-debugging>`_.
1038 equal: Optional[AssignEqual] = None
1039
1040 def _validate(self) -> None:
1041 if self.conversion is not None and self.conversion not in ("s", "r", "a"):
1042 raise CSTValidationError("Invalid t-string conversion.")
1043
1044 def _visit_and_replace_children(
1045 self, visitor: CSTVisitorT
1046 ) -> "TemplatedStringExpression":
1047 format_spec = self.format_spec
1048 return TemplatedStringExpression(
1049 whitespace_before_expression=visit_required(
1050 self,
1051 "whitespace_before_expression",
1052 self.whitespace_before_expression,
1053 visitor,
1054 ),
1055 expression=visit_required(self, "expression", self.expression, visitor),
1056 equal=visit_optional(self, "equal", self.equal, visitor),
1057 whitespace_after_expression=visit_required(
1058 self,
1059 "whitespace_after_expression",
1060 self.whitespace_after_expression,
1061 visitor,
1062 ),
1063 conversion=self.conversion,
1064 format_spec=(
1065 visit_sequence(self, "format_spec", format_spec, visitor)
1066 if format_spec is not None
1067 else None
1068 ),
1069 )
1070
1071 def _codegen_impl(self, state: CodegenState) -> None:
1072 state.add_token("{")
1073 self.whitespace_before_expression._codegen(state)
1074 self.expression._codegen(state)
1075 equal = self.equal
1076 if equal is not None:
1077 equal._codegen(state)
1078 self.whitespace_after_expression._codegen(state)
1079 conversion = self.conversion
1080 if conversion is not None:
1081 state.add_token("!")
1082 state.add_token(conversion)
1083 format_spec = self.format_spec
1084 if format_spec is not None:
1085 state.add_token(":")
1086 for spec in format_spec:
1087 spec._codegen(state)
1088 state.add_token("}")
1089
1090
1091@add_slots
1092@dataclass(frozen=True)
1093class TemplatedString(_BasePrefixedString):
1094 """
1095 An "t-string". Template strings are a generalization of f-strings,
1096 using a t in place of the f prefix. Instead of evaluating to str,
1097 t-strings evaluate to a new type: Template
1098
1099 T-Strings are defined in 'PEP 750'
1100
1101 >>> import libcst as cst
1102 >>> cst.parse_expression('t"ab{cd}ef"')
1103 TemplatedString(
1104 parts=[
1105 TemplatedStringText(
1106 value='ab',
1107 ),
1108 TemplatedStringExpression(
1109 expression=Name(
1110 value='cd',
1111 lpar=[],
1112 rpar=[],
1113 ),
1114 conversion=None,
1115 format_spec=None,
1116 whitespace_before_expression=SimpleWhitespace(
1117 value='',
1118 ),
1119 whitespace_after_expression=SimpleWhitespace(
1120 value='',
1121 ),
1122 equal=None,
1123 ),
1124 TemplatedStringText(
1125 value='ef',
1126 ),
1127 ],
1128 start='t"',
1129 end='"',
1130 lpar=[],
1131 rpar=[],
1132 )
1133 >>>
1134 """
1135
1136 #: A templated string is composed as a series of :class:`TemplatedStringText` and
1137 #: :class:`TemplatedStringExpression` parts.
1138 parts: Sequence[BaseTemplatedStringContent]
1139
1140 #: The string prefix and the leading quote, such as ``t"``, ``T'``, ``tr"``, or
1141 #: ``t"""``.
1142 start: str = 't"'
1143
1144 #: The trailing quote. This must match the type of quote used in ``start``.
1145 end: Literal['"', "'", '"""', "'''"] = '"'
1146
1147 lpar: Sequence[LeftParen] = ()
1148 #: Sequence of parenthesis for precidence dictation.
1149 rpar: Sequence[RightParen] = ()
1150
1151 def _validate(self) -> None:
1152 super(_BasePrefixedString, self)._validate()
1153
1154 # Validate any prefix
1155 prefix = self.prefix
1156 if prefix not in ("t", "tr", "rt"):
1157 raise CSTValidationError("Invalid t-string prefix.")
1158
1159 # Validate wrapping quotes
1160 starttoken = self.start[len(prefix) :]
1161 if starttoken != self.end:
1162 raise CSTValidationError("t-string must have matching enclosing quotes.")
1163
1164 # Validate valid wrapping quote usage
1165 if starttoken not in ('"', "'", '"""', "'''"):
1166 raise CSTValidationError("Invalid t-string enclosing quotes.")
1167
1168 @property
1169 def prefix(self) -> str:
1170 """
1171 Returns the string's prefix, if any exists. The prefix can be ``t``,
1172 ``tr``, or ``rt``.
1173 """
1174
1175 prefix = ""
1176 for c in self.start:
1177 if c in ['"', "'"]:
1178 break
1179 prefix += c
1180 return prefix.lower()
1181
1182 @property
1183 def quote(self) -> StringQuoteLiteral:
1184 """
1185 Returns the quotation used to denote the string. Can be either ``'``,
1186 ``"``, ``'''`` or ``\"\"\"``.
1187 """
1188
1189 return self.end
1190
1191 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "TemplatedString":
1192 return TemplatedString(
1193 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
1194 start=self.start,
1195 parts=visit_sequence(self, "parts", self.parts, visitor),
1196 end=self.end,
1197 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
1198 )
1199
1200 def _codegen_impl(self, state: CodegenState) -> None:
1201 with self._parenthesize(state):
1202 state.add_token(self.start)
1203 for part in self.parts:
1204 part._codegen(state)
1205 state.add_token(self.end)
1206
1207
1208@add_slots
1209@dataclass(frozen=True)
1210class ConcatenatedString(BaseString):
1211 """
1212 Represents an implicitly concatenated string, such as::
1213
1214 "abc" "def" == "abcdef"
1215
1216 .. warning::
1217 This is different from two strings joined in a :class:`BinaryOperation` with an
1218 :class:`Add` operator, and is `sometimes viewed as an antifeature of Python
1219 <https://lwn.net/Articles/551426/>`_.
1220 """
1221
1222 #: String on the left of the concatenation.
1223 left: Union[SimpleString, FormattedString]
1224
1225 #: String on the right of the concatenation.
1226 right: Union[SimpleString, FormattedString, "ConcatenatedString"]
1227
1228 lpar: Sequence[LeftParen] = ()
1229 #: Sequence of parenthesis for precidence dictation.
1230 rpar: Sequence[RightParen] = ()
1231
1232 #: Whitespace between the ``left`` and ``right`` substrings.
1233 whitespace_between: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
1234
1235 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
1236 if super(ConcatenatedString, self)._safe_to_use_with_word_operator(position):
1237 # if we have parenthesis, we're safe.
1238 return True
1239 return self._check_left_right_word_concatenation_safety(
1240 position, self.left, self.right
1241 )
1242
1243 def _validate(self) -> None:
1244 super(ConcatenatedString, self)._validate()
1245
1246 # Strings that are concatenated cannot have parens.
1247 if bool(self.left.lpar) or bool(self.left.rpar):
1248 raise CSTValidationError("Cannot concatenate parenthesized strings.")
1249 if bool(self.right.lpar) or bool(self.right.rpar):
1250 raise CSTValidationError("Cannot concatenate parenthesized strings.")
1251
1252 # Cannot concatenate str and bytes
1253 leftbytes = "b" in self.left.prefix
1254 right = self.right
1255 if isinstance(right, ConcatenatedString):
1256 rightbytes = "b" in right.left.prefix
1257 elif isinstance(right, SimpleString):
1258 rightbytes = "b" in right.prefix
1259 elif isinstance(right, FormattedString):
1260 rightbytes = "b" in right.prefix
1261 else:
1262 raise CSTLogicError("Logic error!")
1263 if leftbytes != rightbytes:
1264 raise CSTValidationError("Cannot concatenate string and bytes.")
1265
1266 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "ConcatenatedString":
1267 return ConcatenatedString(
1268 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
1269 left=visit_required(self, "left", self.left, visitor),
1270 whitespace_between=visit_required(
1271 self, "whitespace_between", self.whitespace_between, visitor
1272 ),
1273 right=visit_required(self, "right", self.right, visitor),
1274 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
1275 )
1276
1277 def _codegen_impl(self, state: CodegenState) -> None:
1278 with self._parenthesize(state):
1279 self.left._codegen(state)
1280 self.whitespace_between._codegen(state)
1281 self.right._codegen(state)
1282
1283 @property
1284 def evaluated_value(self) -> Union[str, bytes, None]:
1285 """
1286 Return an :func:`ast.literal_eval` evaluated str of recursively concatenated :py:attr:`left` and :py:attr:`right`
1287 if and only if both :py:attr:`left` and :py:attr:`right` are composed by :class:`SimpleString` or :class:`ConcatenatedString`
1288 (:class:`FormattedString` cannot be evaluated).
1289 """
1290 left = self.left
1291 right = self.right
1292 if isinstance(left, FormattedString) or isinstance(right, FormattedString):
1293 return None
1294 left_val = left.evaluated_value
1295 right_val = right.evaluated_value
1296 if right_val is None:
1297 return None
1298 if isinstance(left_val, bytes) and isinstance(right_val, bytes):
1299 return left_val + right_val
1300 if isinstance(left_val, str) and isinstance(right_val, str):
1301 return left_val + right_val
1302 return None
1303
1304
1305@add_slots
1306@dataclass(frozen=True)
1307class ComparisonTarget(CSTNode):
1308 """
1309 A target for a :class:`Comparison`. Owns the comparison operator and the value to
1310 the right of the operator.
1311 """
1312
1313 #: A comparison operator such as ``<``, ``>=``, ``==``, ``is``, or ``in``.
1314 operator: BaseCompOp
1315
1316 #: The right hand side of the comparison operation.
1317 comparator: BaseExpression
1318
1319 def _validate(self) -> None:
1320 # Validate operator spacing rules
1321 operator = self.operator
1322 if (
1323 isinstance(operator, (In, NotIn, Is, IsNot))
1324 and operator.whitespace_after.empty
1325 and not self.comparator._safe_to_use_with_word_operator(
1326 ExpressionPosition.RIGHT
1327 )
1328 ):
1329 raise CSTValidationError(
1330 "Must have at least one space around comparison operator."
1331 )
1332
1333 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "ComparisonTarget":
1334 return ComparisonTarget(
1335 operator=visit_required(self, "operator", self.operator, visitor),
1336 comparator=visit_required(self, "comparator", self.comparator, visitor),
1337 )
1338
1339 def _codegen_impl(self, state: CodegenState) -> None:
1340 self.operator._codegen(state)
1341 self.comparator._codegen(state)
1342
1343
1344@add_slots
1345@dataclass(frozen=True)
1346class Comparison(BaseExpression):
1347 """
1348 A comparison between multiple values such as ``x < y``, ``x < y < z``, or
1349 ``x in [y, z]``. These comparisions typically result in boolean values.
1350
1351 Unlike :class:`BinaryOperation` and :class:`BooleanOperation`, comparisons are not
1352 restricted to a left and right child. Instead they can contain an arbitrary number
1353 of :class:`ComparisonTarget` children.
1354
1355 ``x < y < z`` is not equivalent to ``(x < y) < z`` or ``x < (y < z)``. Instead,
1356 it's roughly equivalent to ``x < y and y < z``.
1357
1358 For more details, see `Python's documentation on comparisons
1359 <https://docs.python.org/3/reference/expressions.html#comparisons>`_.
1360
1361 ::
1362
1363 # x < y < z
1364
1365 Comparison(
1366 Name("x"),
1367 [
1368 ComparisonTarget(LessThan(), Name("y")),
1369 ComparisonTarget(LessThan(), Name("z")),
1370 ],
1371 )
1372 """
1373
1374 #: The first value in the full sequence of values to compare. This value will be
1375 #: compared against the first value in ``comparisions``.
1376 left: BaseExpression
1377
1378 #: Pairs of :class:`BaseCompOp` operators and expression values to compare. These
1379 #: come after ``left``. Each value is compared against the value before and after
1380 #: itself in the sequence.
1381 comparisons: Sequence[ComparisonTarget]
1382
1383 lpar: Sequence[LeftParen] = ()
1384 #: Sequence of parenthesis for precedence dictation.
1385 rpar: Sequence[RightParen] = ()
1386
1387 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
1388 if super(Comparison, self)._safe_to_use_with_word_operator(position):
1389 # we have parenthesis
1390 return True
1391 return self._check_left_right_word_concatenation_safety(
1392 position, self.left, self.comparisons[-1].comparator
1393 )
1394
1395 def _validate(self) -> None:
1396 # Perform any validation on base type
1397 super(Comparison, self)._validate()
1398
1399 if len(self.comparisons) == 0:
1400 raise CSTValidationError("Must have at least one ComparisonTarget.")
1401
1402 # Validate operator spacing rules
1403 previous_comparator = self.left
1404 for target in self.comparisons:
1405 operator = target.operator
1406 if (
1407 isinstance(operator, (In, NotIn, Is, IsNot))
1408 and operator.whitespace_before.empty
1409 and not previous_comparator._safe_to_use_with_word_operator(
1410 ExpressionPosition.LEFT
1411 )
1412 ):
1413 raise CSTValidationError(
1414 "Must have at least one space around comparison operator."
1415 )
1416 previous_comparator = target.comparator
1417
1418 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Comparison":
1419 return Comparison(
1420 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
1421 left=visit_required(self, "left", self.left, visitor),
1422 comparisons=visit_sequence(self, "comparisons", self.comparisons, visitor),
1423 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
1424 )
1425
1426 def _codegen_impl(self, state: CodegenState) -> None:
1427 with self._parenthesize(state):
1428 self.left._codegen(state)
1429 for comp in self.comparisons:
1430 comp._codegen(state)
1431
1432
1433@add_slots
1434@dataclass(frozen=True)
1435class UnaryOperation(BaseExpression):
1436 """
1437 Any generic unary expression, such as ``not x`` or ``-x``. :class:`UnaryOperation`
1438 nodes apply a :class:`BaseUnaryOp` to an expression.
1439 """
1440
1441 #: The unary operator that applies some operation (e.g. negation) to the
1442 #: ``expression``.
1443 operator: BaseUnaryOp
1444
1445 #: The expression that should be transformed (e.g. negated) by the operator to
1446 #: create a new value.
1447 expression: BaseExpression
1448
1449 lpar: Sequence[LeftParen] = ()
1450 #: Sequence of parenthesis for precedence dictation.
1451 rpar: Sequence[RightParen] = ()
1452
1453 def _validate(self) -> None:
1454 # Perform any validation on base type
1455 super(UnaryOperation, self)._validate()
1456
1457 if (
1458 isinstance(self.operator, Not)
1459 and self.operator.whitespace_after.empty
1460 and not self.expression._safe_to_use_with_word_operator(
1461 ExpressionPosition.RIGHT
1462 )
1463 ):
1464 raise CSTValidationError("Must have at least one space after not operator.")
1465
1466 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "UnaryOperation":
1467 return UnaryOperation(
1468 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
1469 operator=visit_required(self, "operator", self.operator, visitor),
1470 expression=visit_required(self, "expression", self.expression, visitor),
1471 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
1472 )
1473
1474 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
1475 """
1476 As long as we aren't comprised of the Not unary operator, we are safe to use
1477 without space.
1478 """
1479 if super(UnaryOperation, self)._safe_to_use_with_word_operator(position):
1480 return True
1481 if position == ExpressionPosition.RIGHT:
1482 return not isinstance(self.operator, Not)
1483 if position == ExpressionPosition.LEFT:
1484 return self.expression._safe_to_use_with_word_operator(
1485 ExpressionPosition.LEFT
1486 )
1487 return False
1488
1489 def _codegen_impl(self, state: CodegenState) -> None:
1490 with self._parenthesize(state):
1491 self.operator._codegen(state)
1492 self.expression._codegen(state)
1493
1494
1495@add_slots
1496@dataclass(frozen=True)
1497class BinaryOperation(BaseExpression):
1498 """
1499 An operation that combines two expression such as ``x << y`` or ``y + z``.
1500 :class:`BinaryOperation` nodes apply a :class:`BaseBinaryOp` to an expression.
1501
1502 Binary operations do not include operations performed with :class:`BaseBooleanOp`
1503 nodes, such as ``and`` or ``or``. Instead, those operations are provided by
1504 :class:`BooleanOperation`.
1505
1506 It also does not include support for comparision operators performed with
1507 :class:`BaseCompOp`, such as ``<``, ``>=``, ``==``, ``is``, or ``in``. Instead,
1508 those operations are provided by :class:`Comparison`.
1509 """
1510
1511 #: The left hand side of the operation.
1512 left: BaseExpression
1513
1514 #: The actual operator such as ``<<`` or ``+`` that combines the ``left`` and
1515 #: ``right`` expressions.
1516 operator: BaseBinaryOp
1517
1518 #: The right hand side of the operation.
1519 right: BaseExpression
1520
1521 lpar: Sequence[LeftParen] = ()
1522 #: Sequence of parenthesis for precedence dictation.
1523 rpar: Sequence[RightParen] = ()
1524
1525 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "BinaryOperation":
1526 return BinaryOperation(
1527 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
1528 left=visit_required(self, "left", self.left, visitor),
1529 operator=visit_required(self, "operator", self.operator, visitor),
1530 right=visit_required(self, "right", self.right, visitor),
1531 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
1532 )
1533
1534 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
1535 if super(BinaryOperation, self)._safe_to_use_with_word_operator(position):
1536 return True
1537 return self._check_left_right_word_concatenation_safety(
1538 position, self.left, self.right
1539 )
1540
1541 def _codegen_impl(self, state: CodegenState) -> None:
1542 with self._parenthesize(state):
1543 self.left._codegen(state)
1544 self.operator._codegen(state)
1545 self.right._codegen(state)
1546
1547
1548@add_slots
1549@dataclass(frozen=True)
1550class BooleanOperation(BaseExpression):
1551 """
1552 An operation that combines two booleans such as ``x or y`` or ``z and w``
1553 :class:`BooleanOperation` nodes apply a :class:`BaseBooleanOp` to an expression.
1554
1555 Boolean operations do not include operations performed with :class:`BaseBinaryOp`
1556 nodes, such as ``+`` or ``<<``. Instead, those operations are provided by
1557 :class:`BinaryOperation`.
1558
1559 It also does not include support for comparision operators performed with
1560 :class:`BaseCompOp`, such as ``<``, ``>=``, ``==``, ``is``, or ``in``. Instead,
1561 those operations are provided by :class:`Comparison`.
1562 """
1563
1564 #: The left hand side of the operation.
1565 left: BaseExpression
1566
1567 #: The actual operator such as ``and`` or ``or`` that combines the ``left`` and
1568 #: ``right`` expressions.
1569 operator: BaseBooleanOp
1570
1571 #: The right hand side of the operation.
1572 right: BaseExpression
1573
1574 lpar: Sequence[LeftParen] = ()
1575 #: Sequence of parenthesis for precedence dictation.
1576 rpar: Sequence[RightParen] = ()
1577
1578 def _validate(self) -> None:
1579 # Paren validation and such
1580 super(BooleanOperation, self)._validate()
1581 # Validate spacing rules
1582 if (
1583 self.operator.whitespace_before.empty
1584 and not self.left._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
1585 ):
1586 raise CSTValidationError(
1587 "Must have at least one space around boolean operator."
1588 )
1589 if (
1590 self.operator.whitespace_after.empty
1591 and not self.right._safe_to_use_with_word_operator(ExpressionPosition.RIGHT)
1592 ):
1593 raise CSTValidationError(
1594 "Must have at least one space around boolean operator."
1595 )
1596
1597 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "BooleanOperation":
1598 return BooleanOperation(
1599 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
1600 left=visit_required(self, "left", self.left, visitor),
1601 operator=visit_required(self, "operator", self.operator, visitor),
1602 right=visit_required(self, "right", self.right, visitor),
1603 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
1604 )
1605
1606 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
1607 if super(BooleanOperation, self)._safe_to_use_with_word_operator(position):
1608 return True
1609 return self._check_left_right_word_concatenation_safety(
1610 position, self.left, self.right
1611 )
1612
1613 def _codegen_impl(self, state: CodegenState) -> None:
1614 with self._parenthesize(state):
1615 self.left._codegen(state)
1616 self.operator._codegen(state)
1617 self.right._codegen(state)
1618
1619
1620@add_slots
1621@dataclass(frozen=True)
1622class Attribute(BaseAssignTargetExpression, BaseDelTargetExpression):
1623 """
1624 An attribute reference, such as ``x.y``.
1625
1626 Note that in the case of ``x.y.z``, the outer attribute will have an attr of ``z``
1627 and the value will be another :class:`Attribute` referencing the ``y`` attribute on
1628 ``x``::
1629
1630 Attribute(
1631 value=Attribute(
1632 value=Name("x")
1633 attr=Name("y")
1634 ),
1635 attr=Name("z"),
1636 )
1637 """
1638
1639 #: An expression which, when evaluated, will produce an object with ``attr`` as an
1640 #: attribute.
1641 value: BaseExpression
1642
1643 #: The name of the attribute being accessed on the ``value`` object.
1644 attr: Name
1645
1646 #: A separating dot. If there's whitespace between the ``value`` and ``attr``, this
1647 #: dot owns it.
1648 dot: Dot = Dot()
1649
1650 lpar: Sequence[LeftParen] = ()
1651 #: Sequence of parenthesis for precedence dictation.
1652 rpar: Sequence[RightParen] = ()
1653
1654 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Attribute":
1655 return Attribute(
1656 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
1657 value=visit_required(self, "value", self.value, visitor),
1658 dot=visit_required(self, "dot", self.dot, visitor),
1659 attr=visit_required(self, "attr", self.attr, visitor),
1660 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
1661 )
1662
1663 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
1664 if super(Attribute, self)._safe_to_use_with_word_operator(position):
1665 return True
1666 return self._check_left_right_word_concatenation_safety(
1667 position, self.value, self.attr
1668 )
1669
1670 def _codegen_impl(self, state: CodegenState) -> None:
1671 with self._parenthesize(state):
1672 self.value._codegen(state)
1673 self.dot._codegen(state)
1674 self.attr._codegen(state)
1675
1676
1677class BaseSlice(CSTNode, ABC):
1678 """
1679 Any slice type that can slot into a :class:`SubscriptElement`.
1680 This node is purely for typing.
1681 """
1682
1683 __slots__ = ()
1684
1685
1686@add_slots
1687@dataclass(frozen=True)
1688class Index(BaseSlice):
1689 """
1690 Any index as passed to a :class:`Subscript`. In ``x[2]``, this would be the ``2``
1691 value.
1692 """
1693
1694 #: The index value itself.
1695 value: BaseExpression
1696
1697 #: An optional string with an asterisk appearing before the name. This is
1698 #: expanded into variable number of positional arguments. See PEP-646
1699 star: Optional[Literal["*"]] = None
1700
1701 #: Whitespace after the ``star`` (if it exists), but before the ``value``.
1702 whitespace_after_star: Optional[BaseParenthesizableWhitespace] = None
1703
1704 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Index":
1705 return Index(
1706 star=self.star,
1707 whitespace_after_star=visit_optional(
1708 self, "whitespace_after_star", self.whitespace_after_star, visitor
1709 ),
1710 value=visit_required(self, "value", self.value, visitor),
1711 )
1712
1713 def _codegen_impl(self, state: CodegenState) -> None:
1714 star = self.star
1715 if star is not None:
1716 state.add_token(star)
1717 ws = self.whitespace_after_star
1718 if ws is not None:
1719 ws._codegen(state)
1720 self.value._codegen(state)
1721
1722
1723@add_slots
1724@dataclass(frozen=True)
1725class Slice(BaseSlice):
1726 """
1727 Any slice operation in a :class:`Subscript`, such as ``1:``, ``2:3:4``, etc.
1728
1729 Note that the grammar does NOT allow parenthesis around a slice so they are not
1730 supported here.
1731 """
1732
1733 #: The lower bound in the slice, if present
1734 lower: Optional[BaseExpression]
1735
1736 #: The upper bound in the slice, if present
1737 upper: Optional[BaseExpression]
1738
1739 #: The step in the slice, if present
1740 step: Optional[BaseExpression] = None
1741
1742 #: The first slice operator
1743 first_colon: Colon = Colon.field()
1744
1745 #: The second slice operator, usually omitted
1746 second_colon: Union[Colon, MaybeSentinel] = MaybeSentinel.DEFAULT
1747
1748 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Slice":
1749 return Slice(
1750 lower=visit_optional(self, "lower", self.lower, visitor),
1751 first_colon=visit_required(self, "first_colon", self.first_colon, visitor),
1752 upper=visit_optional(self, "upper", self.upper, visitor),
1753 second_colon=visit_sentinel(
1754 self, "second_colon", self.second_colon, visitor
1755 ),
1756 step=visit_optional(self, "step", self.step, visitor),
1757 )
1758
1759 def _codegen_impl(self, state: CodegenState) -> None:
1760 lower = self.lower
1761 if lower is not None:
1762 lower._codegen(state)
1763 self.first_colon._codegen(state)
1764 upper = self.upper
1765 if upper is not None:
1766 upper._codegen(state)
1767 second_colon = self.second_colon
1768 if second_colon is MaybeSentinel.DEFAULT and self.step is not None:
1769 state.add_token(":")
1770 elif isinstance(second_colon, Colon):
1771 second_colon._codegen(state)
1772 step = self.step
1773 if step is not None:
1774 step._codegen(state)
1775
1776
1777@add_slots
1778@dataclass(frozen=True)
1779class SubscriptElement(CSTNode):
1780 """
1781 Part of a sequence of slices in a :class:`Subscript`, such as ``1:2, 3``. This is
1782 not used in Python's standard library, but it is used in some third-party
1783 libraries. For example, `NumPy uses it to select values and ranges from
1784 multi-dimensional arrays
1785 <https://docs.scipy.org/doc/numpy-1.10.1/user/basics.indexing.html>`_.
1786 """
1787
1788 #: A slice or index that is part of a subscript.
1789 slice: BaseSlice
1790
1791 #: A separating comma, with any whitespace it owns.
1792 comma: Union[Comma, MaybeSentinel] = MaybeSentinel.DEFAULT
1793
1794 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "SubscriptElement":
1795 return SubscriptElement(
1796 slice=visit_required(self, "slice", self.slice, visitor),
1797 comma=visit_sentinel(self, "comma", self.comma, visitor),
1798 )
1799
1800 def _codegen_impl(self, state: CodegenState, default_comma: bool = False) -> None:
1801 with state.record_syntactic_position(self):
1802 self.slice._codegen(state)
1803
1804 comma = self.comma
1805 if comma is MaybeSentinel.DEFAULT and default_comma:
1806 state.add_token(", ")
1807 elif isinstance(comma, Comma):
1808 comma._codegen(state)
1809
1810
1811@add_slots
1812@dataclass(frozen=True)
1813class Subscript(BaseAssignTargetExpression, BaseDelTargetExpression):
1814 """
1815 A indexed subscript reference (:class:`Index`) such as ``x[2]``, a :class:`Slice`
1816 such as ``x[1:-1]``, or an extended slice (:class:`SubscriptElement`) such as ``x[1:2, 3]``.
1817 """
1818
1819 #: The left-hand expression which, when evaluated, will be subscripted, such as
1820 #: ``x`` in ``x[2]``.
1821 value: BaseExpression
1822
1823 #: The :class:`SubscriptElement` to extract from the ``value``.
1824 slice: Sequence[SubscriptElement]
1825
1826 lbracket: LeftSquareBracket = LeftSquareBracket.field()
1827 #: Brackets after the ``value`` surrounding the ``slice``.
1828 rbracket: RightSquareBracket = RightSquareBracket.field()
1829
1830 lpar: Sequence[LeftParen] = ()
1831 #: Sequence of parenthesis for precedence dictation.
1832 rpar: Sequence[RightParen] = ()
1833
1834 #: Whitespace after the ``value``, but before the ``lbracket``.
1835 whitespace_after_value: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
1836
1837 def _validate(self) -> None:
1838 super(Subscript, self)._validate()
1839 # Validate valid commas
1840 if len(self.slice) < 1:
1841 raise CSTValidationError("Cannot have empty SubscriptElement.")
1842
1843 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Subscript":
1844 return Subscript(
1845 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
1846 value=visit_required(self, "value", self.value, visitor),
1847 whitespace_after_value=visit_required(
1848 self, "whitespace_after_value", self.whitespace_after_value, visitor
1849 ),
1850 lbracket=visit_required(self, "lbracket", self.lbracket, visitor),
1851 slice=visit_sequence(self, "slice", self.slice, visitor),
1852 rbracket=visit_required(self, "rbracket", self.rbracket, visitor),
1853 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
1854 )
1855
1856 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
1857 if position == ExpressionPosition.LEFT:
1858 return True
1859 if super(Subscript, self)._safe_to_use_with_word_operator(position):
1860 return True
1861 if position == ExpressionPosition.RIGHT:
1862 return self.value._safe_to_use_with_word_operator(ExpressionPosition.RIGHT)
1863 return False
1864
1865 def _codegen_impl(self, state: CodegenState) -> None:
1866 with self._parenthesize(state):
1867 self.value._codegen(state)
1868 self.whitespace_after_value._codegen(state)
1869 self.lbracket._codegen(state)
1870 lastslice = len(self.slice) - 1
1871 for i, slice in enumerate(self.slice):
1872 slice._codegen(state, default_comma=(i != lastslice))
1873 self.rbracket._codegen(state)
1874
1875
1876@add_slots
1877@dataclass(frozen=True)
1878class Annotation(CSTNode):
1879 """
1880 An annotation for a function (`PEP 3107`_) or on a variable (`PEP 526`_). Typically
1881 these are used in the context of type hints (`PEP 484`_), such as::
1882
1883 # a variable with a type
1884 good_ideas: List[str] = []
1885
1886 # a function with type annotations
1887 def concat(substrings: Sequence[str]) -> str:
1888 ...
1889
1890 .. _PEP 3107: https://www.python.org/dev/peps/pep-3107/
1891 .. _PEP 526: https://www.python.org/dev/peps/pep-0526/
1892 .. _PEP 484: https://www.python.org/dev/peps/pep-0484/
1893 """
1894
1895 #: The annotation's value itself. This is the part of the annotation after the
1896 #: colon or arrow.
1897 annotation: BaseExpression
1898
1899 whitespace_before_indicator: Union[BaseParenthesizableWhitespace, MaybeSentinel] = (
1900 MaybeSentinel.DEFAULT
1901 )
1902 whitespace_after_indicator: BaseParenthesizableWhitespace = SimpleWhitespace.field(
1903 " "
1904 )
1905
1906 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Annotation":
1907 return Annotation(
1908 whitespace_before_indicator=visit_sentinel(
1909 self,
1910 "whitespace_before_indicator",
1911 self.whitespace_before_indicator,
1912 visitor,
1913 ),
1914 whitespace_after_indicator=visit_required(
1915 self,
1916 "whitespace_after_indicator",
1917 self.whitespace_after_indicator,
1918 visitor,
1919 ),
1920 annotation=visit_required(self, "annotation", self.annotation, visitor),
1921 )
1922
1923 def _codegen_impl(
1924 self, state: CodegenState, default_indicator: Optional[str] = None
1925 ) -> None:
1926 # First, figure out the indicator which tells us default whitespace.
1927 if default_indicator is None:
1928 raise CSTCodegenError(
1929 "Must specify a concrete default_indicator if default used on indicator."
1930 )
1931
1932 # Now, output the whitespace
1933 whitespace_before_indicator = self.whitespace_before_indicator
1934 if isinstance(whitespace_before_indicator, BaseParenthesizableWhitespace):
1935 whitespace_before_indicator._codegen(state)
1936 elif isinstance(whitespace_before_indicator, MaybeSentinel):
1937 if default_indicator == "->":
1938 state.add_token(" ")
1939 else:
1940 raise CSTLogicError("Logic error!")
1941
1942 # Now, output the indicator and the rest of the annotation
1943 state.add_token(default_indicator)
1944 self.whitespace_after_indicator._codegen(state)
1945
1946 with state.record_syntactic_position(self):
1947 self.annotation._codegen(state)
1948
1949
1950@add_slots
1951@dataclass(frozen=True)
1952class ParamStar(CSTNode):
1953 """
1954 A sentinel indicator on a :class:`Parameters` list to denote that the subsequent
1955 params are keyword-only args.
1956
1957 This syntax is described in `PEP 3102`_.
1958
1959 .. _PEP 3102: https://www.python.org/dev/peps/pep-3102/#specification
1960 """
1961
1962 # Comma that comes after the star.
1963 comma: Comma = Comma.field(whitespace_after=SimpleWhitespace(" "))
1964
1965 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "ParamStar":
1966 return ParamStar(comma=visit_required(self, "comma", self.comma, visitor))
1967
1968 def _codegen_impl(self, state: CodegenState) -> None:
1969 state.add_token("*")
1970 self.comma._codegen(state)
1971
1972
1973@add_slots
1974@dataclass(frozen=True)
1975class ParamSlash(CSTNode):
1976 """
1977 A sentinel indicator on a :class:`Parameters` list to denote that the previous
1978 params are positional-only args.
1979
1980 This syntax is described in `PEP 570`_.
1981
1982 .. _PEP 570: https://www.python.org/dev/peps/pep-0570/#specification
1983 """
1984
1985 #: Optional comma that comes after the slash. This comma doesn't own the whitespace
1986 #: between ``/`` and ``,``.
1987 comma: Union[Comma, MaybeSentinel] = MaybeSentinel.DEFAULT
1988
1989 #: Whitespace after the ``/`` character. This is captured here in case there is a
1990 #: comma.
1991 whitespace_after: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
1992
1993 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "ParamSlash":
1994 return ParamSlash(
1995 comma=visit_sentinel(self, "comma", self.comma, visitor),
1996 whitespace_after=visit_required(
1997 self, "whitespace_after", self.whitespace_after, visitor
1998 ),
1999 )
2000
2001 def _codegen_impl(self, state: CodegenState, default_comma: bool = False) -> None:
2002 state.add_token("/")
2003
2004 self.whitespace_after._codegen(state)
2005 comma = self.comma
2006 if comma is MaybeSentinel.DEFAULT and default_comma:
2007 state.add_token(", ")
2008 elif isinstance(comma, Comma):
2009 comma._codegen(state)
2010
2011
2012@add_slots
2013@dataclass(frozen=True)
2014class Param(CSTNode):
2015 """
2016 A positional or keyword argument in a :class:`Parameters` list. May contain an
2017 :class:`Annotation` and, in some cases, a ``default``.
2018 """
2019
2020 #: The parameter name itself.
2021 name: Name
2022
2023 #: Any optional :class:`Annotation`. These annotations are usually used as type
2024 #: hints.
2025 annotation: Optional[Annotation] = None
2026
2027 #: The equal sign used to denote assignment if there is a default.
2028 equal: Union[AssignEqual, MaybeSentinel] = MaybeSentinel.DEFAULT
2029
2030 #: Any optional default value, used when the argument is not supplied.
2031 default: Optional[BaseExpression] = None
2032
2033 #: A trailing comma. If one is not provided, :class:`MaybeSentinel` will be
2034 #: replaced with a comma only if a comma is required.
2035 comma: Union[Comma, MaybeSentinel] = MaybeSentinel.DEFAULT
2036
2037 #: Zero, one, or two asterisks appearing before name for :class:`Param`'s
2038 #: ``star_arg`` and ``star_kwarg``.
2039 star: Union[str, MaybeSentinel] = MaybeSentinel.DEFAULT
2040
2041 #: The whitespace before ``name``. It will appear after ``star`` when a star
2042 #: exists.
2043 whitespace_after_star: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
2044
2045 #: The whitespace after this entire node.
2046 whitespace_after_param: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
2047
2048 def _validate(self) -> None:
2049 if self.default is None and isinstance(self.equal, AssignEqual):
2050 raise CSTValidationError(
2051 "Must have a default when specifying an AssignEqual."
2052 )
2053 if isinstance(self.star, str) and self.star not in ("", "*", "**"):
2054 raise CSTValidationError("Must specify either '', '*' or '**' for star.")
2055
2056 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Param":
2057 return Param(
2058 star=self.star,
2059 whitespace_after_star=visit_required(
2060 self, "whitespace_after_star", self.whitespace_after_star, visitor
2061 ),
2062 name=visit_required(self, "name", self.name, visitor),
2063 annotation=visit_optional(self, "annotation", self.annotation, visitor),
2064 equal=visit_sentinel(self, "equal", self.equal, visitor),
2065 default=visit_optional(self, "default", self.default, visitor),
2066 comma=visit_sentinel(self, "comma", self.comma, visitor),
2067 whitespace_after_param=visit_required(
2068 self, "whitespace_after_param", self.whitespace_after_param, visitor
2069 ),
2070 )
2071
2072 def _codegen_impl(
2073 self,
2074 state: CodegenState,
2075 default_star: Optional[str] = None,
2076 default_comma: bool = False,
2077 ) -> None:
2078 with state.record_syntactic_position(self):
2079 star = self.star
2080 if isinstance(star, MaybeSentinel):
2081 if default_star is None:
2082 raise CSTCodegenError(
2083 "Must specify a concrete default_star if default used on star."
2084 )
2085 star = default_star
2086 if isinstance(star, str):
2087 state.add_token(star)
2088 self.whitespace_after_star._codegen(state)
2089 self.name._codegen(state)
2090
2091 annotation = self.annotation
2092 if annotation is not None:
2093 annotation._codegen(state, default_indicator=":")
2094 equal = self.equal
2095 if equal is MaybeSentinel.DEFAULT and self.default is not None:
2096 state.add_token(" = ")
2097 elif isinstance(equal, AssignEqual):
2098 equal._codegen(state)
2099 default = self.default
2100 if default is not None:
2101 default._codegen(state)
2102 comma = self.comma
2103 if comma is MaybeSentinel.DEFAULT and default_comma:
2104 state.add_token(", ")
2105 elif isinstance(comma, Comma):
2106 comma._codegen(state)
2107
2108 self.whitespace_after_param._codegen(state)
2109
2110
2111@add_slots
2112@dataclass(frozen=True)
2113class Parameters(CSTNode):
2114 """
2115 A function or lambda parameter list.
2116 """
2117
2118 #: Positional parameters, with or without defaults. Positional parameters
2119 #: with defaults must all be after those without defaults.
2120 params: Sequence[Param] = ()
2121
2122 # Optional parameter that captures unspecified positional arguments or a sentinel
2123 # star that dictates parameters following are kwonly args.
2124 star_arg: Union[Param, ParamStar, MaybeSentinel] = MaybeSentinel.DEFAULT
2125
2126 #: Keyword-only params that may or may not have defaults.
2127 kwonly_params: Sequence[Param] = ()
2128
2129 #: Optional parameter that captures unspecified kwargs.
2130 star_kwarg: Optional[Param] = None
2131
2132 #: Positional-only parameters, with or without defaults. Positional-only
2133 #: parameters with defaults must all be after those without defaults.
2134 posonly_params: Sequence[Param] = ()
2135
2136 #: Optional sentinel that dictates parameters preceeding are positional-only
2137 #: args.
2138 posonly_ind: Union[ParamSlash, MaybeSentinel] = MaybeSentinel.DEFAULT
2139
2140 def _validate_stars_sequence(self, vals: Sequence[Param], *, section: str) -> None:
2141 if len(vals) == 0:
2142 return
2143 for val in vals:
2144 if isinstance(val.star, str) and val.star != "":
2145 raise CSTValidationError(
2146 f"Expecting a star prefix of '' for {section} Param."
2147 )
2148
2149 def _validate_posonly_ind(self) -> None:
2150 if isinstance(self.posonly_ind, ParamSlash) and len(self.posonly_params) == 0:
2151 raise CSTValidationError(
2152 "Must have at least one posonly param if ParamSlash is used."
2153 )
2154
2155 def _validate_kwonly_star(self) -> None:
2156 if isinstance(self.star_arg, ParamStar) and len(self.kwonly_params) == 0:
2157 raise CSTValidationError(
2158 "Must have at least one kwonly param if ParamStar is used."
2159 )
2160
2161 def _validate_defaults(self) -> None:
2162 seen_default = False
2163 # pyre-fixme[60]: Concatenation not yet support for multiple variadic
2164 # tuples: `*self.posonly_params, *self.params`.
2165 for param in (*self.posonly_params, *self.params):
2166 if param.default:
2167 # Mark that we've moved onto defaults
2168 if not seen_default:
2169 seen_default = True
2170 else:
2171 if seen_default:
2172 # We accidentally included a non-default after a default arg!
2173 raise CSTValidationError(
2174 "Cannot have param without defaults following a param with defaults."
2175 )
2176 star_arg = self.star_arg
2177 if isinstance(star_arg, Param) and star_arg.default is not None:
2178 raise CSTValidationError("Cannot have default for star_arg.")
2179 star_kwarg = self.star_kwarg
2180 if star_kwarg is not None and star_kwarg.default is not None:
2181 raise CSTValidationError("Cannot have default for star_kwarg.")
2182
2183 def _validate_stars(self) -> None:
2184 if len(self.params) > 0:
2185 self._validate_stars_sequence(self.params, section="params")
2186 if len(self.posonly_params) > 0:
2187 self._validate_stars_sequence(self.posonly_params, section="posonly_params")
2188 star_arg = self.star_arg
2189 if (
2190 isinstance(star_arg, Param)
2191 and isinstance(star_arg.star, str)
2192 and star_arg.star != "*"
2193 ):
2194 raise CSTValidationError(
2195 "Expecting a star prefix of '*' for star_arg Param."
2196 )
2197 if len(self.kwonly_params) > 0:
2198 self._validate_stars_sequence(self.kwonly_params, section="kwonly_params")
2199 star_kwarg = self.star_kwarg
2200 if (
2201 star_kwarg is not None
2202 and isinstance(star_kwarg.star, str)
2203 and star_kwarg.star != "**"
2204 ):
2205 raise CSTValidationError(
2206 "Expecting a star prefix of '**' for star_kwarg Param."
2207 )
2208
2209 def _validate(self) -> None:
2210 # Validate posonly_params slash placement semantics.
2211 self._validate_posonly_ind()
2212 # Validate kwonly_param star placement semantics.
2213 self._validate_kwonly_star()
2214 # Validate defaults semantics for params and star_arg/star_kwarg.
2215 self._validate_defaults()
2216 # Validate that we don't have random stars on non star_kwarg.
2217 self._validate_stars()
2218
2219 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Parameters":
2220 return Parameters(
2221 posonly_params=visit_sequence(
2222 self, "posonly_params", self.posonly_params, visitor
2223 ),
2224 posonly_ind=visit_sentinel(self, "posonly_ind", self.posonly_ind, visitor),
2225 params=visit_sequence(self, "params", self.params, visitor),
2226 star_arg=visit_sentinel(self, "star_arg", self.star_arg, visitor),
2227 kwonly_params=visit_sequence(
2228 self, "kwonly_params", self.kwonly_params, visitor
2229 ),
2230 star_kwarg=visit_optional(self, "star_kwarg", self.star_kwarg, visitor),
2231 )
2232
2233 def _safe_to_join_with_lambda(self) -> bool:
2234 """
2235 Determine if Parameters need a space after the `lambda` keyword. Returns True
2236 iff it's safe to omit the space between `lambda` and these Parameters.
2237
2238 See also `BaseExpression._safe_to_use_with_word_operator`.
2239
2240 For example: `lambda*_: pass`
2241 """
2242 if len(self.posonly_params) != 0:
2243 return False
2244
2245 # posonly_ind can't appear if above condition is false
2246
2247 if len(self.params) > 0 and self.params[0].star not in {"*", "**"}:
2248 return False
2249
2250 return True
2251
2252 def _codegen_impl(self, state: CodegenState) -> None: # noqa: C901
2253 # Compute the star existence first so we can ask about whether
2254 # each element is the last in the list or not.
2255 star_arg = self.star_arg
2256 if isinstance(star_arg, MaybeSentinel):
2257 starincluded = len(self.kwonly_params) > 0
2258 elif isinstance(star_arg, (Param, ParamStar)):
2259 starincluded = True
2260 else:
2261 starincluded = False
2262 # Render out the positional-only params first. They will always have trailing
2263 # commas because in order to have positional-only params, there must be a
2264 # slash afterwards.
2265 for i, param in enumerate(self.posonly_params):
2266 param._codegen(state, default_star="", default_comma=True)
2267 # Render out the positional-only indicator if necessary.
2268 more_values = (
2269 starincluded
2270 or len(self.params) > 0
2271 or len(self.kwonly_params) > 0
2272 or self.star_kwarg is not None
2273 )
2274 posonly_ind = self.posonly_ind
2275 if isinstance(posonly_ind, ParamSlash):
2276 # Its explicitly included, so render the version we have here which
2277 # might have spacing applied to its comma.
2278 posonly_ind._codegen(state, default_comma=more_values)
2279 elif len(self.posonly_params) > 0:
2280 if more_values:
2281 state.add_token("/, ")
2282 else:
2283 state.add_token("/")
2284 # Render out the params next, computing necessary trailing commas.
2285 lastparam = len(self.params) - 1
2286 more_values = (
2287 starincluded or len(self.kwonly_params) > 0 or self.star_kwarg is not None
2288 )
2289 for i, param in enumerate(self.params):
2290 param._codegen(
2291 state, default_star="", default_comma=(i < lastparam or more_values)
2292 )
2293 # Render out optional star sentinel if its explicitly included or
2294 # if we are inferring it from kwonly_params. Otherwise, render out the
2295 # optional star_arg.
2296 if isinstance(star_arg, MaybeSentinel):
2297 if starincluded:
2298 state.add_token("*, ")
2299 elif isinstance(star_arg, Param):
2300 more_values = len(self.kwonly_params) > 0 or self.star_kwarg is not None
2301 star_arg._codegen(state, default_star="*", default_comma=more_values)
2302 elif isinstance(star_arg, ParamStar):
2303 star_arg._codegen(state)
2304 # Render out the kwonly_args next, computing necessary trailing commas.
2305 lastparam = len(self.kwonly_params) - 1
2306 more_values = self.star_kwarg is not None
2307 for i, param in enumerate(self.kwonly_params):
2308 param._codegen(
2309 state, default_star="", default_comma=(i < lastparam or more_values)
2310 )
2311 # Finally, render out any optional star_kwarg
2312 star_kwarg = self.star_kwarg
2313 if star_kwarg is not None:
2314 star_kwarg._codegen(state, default_star="**", default_comma=False)
2315
2316
2317@add_slots
2318@dataclass(frozen=True)
2319class Lambda(BaseExpression):
2320 """
2321 A lambda expression that creates an anonymous function.
2322
2323 ::
2324
2325 Lambda(
2326 params=Parameters([Param(Name("arg"))]),
2327 body=Ellipsis(),
2328 )
2329
2330 Represents the following code::
2331
2332 lambda arg: ...
2333
2334 Named functions statements are provided by :class:`FunctionDef`.
2335 """
2336
2337 #: The arguments to the lambda. This is similar to the arguments on a
2338 #: :class:`FunctionDef`, however lambda arguments are not allowed to have an
2339 #: :class:`Annotation`.
2340 params: Parameters
2341
2342 #: The value that the lambda computes and returns when called.
2343 body: BaseExpression
2344
2345 #: The colon separating the parameters from the body.
2346 colon: Colon = Colon.field(whitespace_after=SimpleWhitespace(" "))
2347
2348 lpar: Sequence[LeftParen] = ()
2349 #: Sequence of parenthesis for precedence dictation.
2350 rpar: Sequence[RightParen] = ()
2351
2352 #: Whitespace after the lambda keyword, but before any argument or the colon.
2353 whitespace_after_lambda: Union[BaseParenthesizableWhitespace, MaybeSentinel] = (
2354 MaybeSentinel.DEFAULT
2355 )
2356
2357 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
2358 if position == ExpressionPosition.LEFT:
2359 return len(self.rpar) > 0 or self.body._safe_to_use_with_word_operator(
2360 position
2361 )
2362 return super()._safe_to_use_with_word_operator(position)
2363
2364 def _validate(self) -> None:
2365 # Validate parents
2366 super(Lambda, self)._validate()
2367 # Sum up all parameters
2368 all_params = [
2369 *self.params.posonly_params,
2370 *self.params.params,
2371 *self.params.kwonly_params,
2372 ]
2373 star_arg = self.params.star_arg
2374 if isinstance(star_arg, Param):
2375 all_params.append(star_arg)
2376 star_kwarg = self.params.star_kwarg
2377 if star_kwarg is not None:
2378 all_params.append(star_kwarg)
2379 # Check for nonzero parameters because several checks care
2380 # about this.
2381 if len(all_params) > 0:
2382 for param in all_params:
2383 if param.annotation is not None:
2384 raise CSTValidationError(
2385 "Lambda params cannot have type annotations."
2386 )
2387 whitespace_after_lambda = self.whitespace_after_lambda
2388 if (
2389 isinstance(whitespace_after_lambda, BaseParenthesizableWhitespace)
2390 and whitespace_after_lambda.empty
2391 and not self.params._safe_to_join_with_lambda()
2392 ):
2393 raise CSTValidationError(
2394 "Must have at least one space after lambda when specifying params"
2395 )
2396
2397 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Lambda":
2398 return Lambda(
2399 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
2400 whitespace_after_lambda=visit_sentinel(
2401 self, "whitespace_after_lambda", self.whitespace_after_lambda, visitor
2402 ),
2403 params=visit_required(self, "params", self.params, visitor),
2404 colon=visit_required(self, "colon", self.colon, visitor),
2405 body=visit_required(self, "body", self.body, visitor),
2406 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
2407 )
2408
2409 def _codegen_impl(self, state: CodegenState) -> None:
2410 with self._parenthesize(state):
2411 state.add_token("lambda")
2412 whitespace_after_lambda = self.whitespace_after_lambda
2413 if isinstance(whitespace_after_lambda, MaybeSentinel):
2414 if not (
2415 len(self.params.posonly_params) == 0
2416 and len(self.params.params) == 0
2417 and not isinstance(self.params.star_arg, Param)
2418 and len(self.params.kwonly_params) == 0
2419 and self.params.star_kwarg is None
2420 ):
2421 # We have one or more params, provide a space
2422 state.add_token(" ")
2423 elif isinstance(whitespace_after_lambda, BaseParenthesizableWhitespace):
2424 whitespace_after_lambda._codegen(state)
2425 self.params._codegen(state)
2426 self.colon._codegen(state)
2427 self.body._codegen(state)
2428
2429
2430@add_slots
2431@dataclass(frozen=True)
2432class Arg(CSTNode):
2433 """
2434 A single argument to a :class:`Call`.
2435
2436 This supports named keyword arguments in the form of ``keyword=value`` and variable
2437 argument expansion using ``*args`` or ``**kwargs`` syntax.
2438 """
2439
2440 #: The argument expression itself, not including a preceding keyword, or any of
2441 #: the surrounding the value, like a comma or asterisks.
2442 value: BaseExpression
2443
2444 #: Optional keyword for the argument.
2445 keyword: Optional[Name] = None
2446
2447 #: The equal sign used to denote assignment if there is a keyword.
2448 equal: Union[AssignEqual, MaybeSentinel] = MaybeSentinel.DEFAULT
2449
2450 #: Any trailing comma.
2451 comma: Union[Comma, MaybeSentinel] = MaybeSentinel.DEFAULT
2452
2453 #: A string with zero, one, or two asterisks appearing before the name. These are
2454 #: expanded into variable number of positional or keyword arguments.
2455 star: Literal["", "*", "**"] = ""
2456
2457 #: Whitespace after the ``star`` (if it exists), but before the ``keyword`` or
2458 #: ``value`` (if no keyword is provided).
2459 whitespace_after_star: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
2460 #: Whitespace after this entire node. The :class:`Comma` node (if it exists) may
2461 #: also store some trailing whitespace.
2462 whitespace_after_arg: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
2463
2464 def _validate(self) -> None:
2465 if self.keyword is None and isinstance(self.equal, AssignEqual):
2466 raise CSTValidationError(
2467 "Must have a keyword when specifying an AssignEqual."
2468 )
2469 if self.star not in ("", "*", "**"):
2470 raise CSTValidationError("Must specify either '', '*' or '**' for star.")
2471 if self.star in ("*", "**") and self.keyword is not None:
2472 raise CSTValidationError("Cannot specify a star and a keyword together.")
2473
2474 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Arg":
2475 return Arg(
2476 star=self.star,
2477 whitespace_after_star=visit_required(
2478 self, "whitespace_after_star", self.whitespace_after_star, visitor
2479 ),
2480 keyword=visit_optional(self, "keyword", self.keyword, visitor),
2481 equal=visit_sentinel(self, "equal", self.equal, visitor),
2482 value=visit_required(self, "value", self.value, visitor),
2483 comma=visit_sentinel(self, "comma", self.comma, visitor),
2484 whitespace_after_arg=visit_required(
2485 self, "whitespace_after_arg", self.whitespace_after_arg, visitor
2486 ),
2487 )
2488
2489 def _codegen_impl(self, state: CodegenState, default_comma: bool = False) -> None:
2490 with state.record_syntactic_position(self):
2491 state.add_token(self.star)
2492 self.whitespace_after_star._codegen(state)
2493 keyword = self.keyword
2494 if keyword is not None:
2495 keyword._codegen(state)
2496 equal = self.equal
2497 if equal is MaybeSentinel.DEFAULT and self.keyword is not None:
2498 state.add_token(" = ")
2499 elif isinstance(equal, AssignEqual):
2500 equal._codegen(state)
2501 self.value._codegen(state)
2502
2503 comma = self.comma
2504 if comma is MaybeSentinel.DEFAULT and default_comma:
2505 state.add_token(", ")
2506 elif isinstance(comma, Comma):
2507 comma._codegen(state)
2508 self.whitespace_after_arg._codegen(state)
2509
2510
2511class _BaseExpressionWithArgs(BaseExpression, ABC):
2512 """
2513 Arguments are complicated enough that we can't represent them easily
2514 in typing. So, we have common validation functions here.
2515 """
2516
2517 __slots__ = ()
2518
2519 #: Sequence of arguments that will be passed to the function call.
2520 args: Sequence[Arg] = ()
2521
2522 def _check_kwargs_or_keywords(self, arg: Arg) -> None:
2523 """
2524 Validates that we only have a mix of "keyword=arg" and "**arg" expansion.
2525 """
2526
2527 if arg.keyword is not None:
2528 # Valid, keyword argument
2529 return None
2530 elif arg.star == "**":
2531 # Valid, kwargs
2532 return None
2533 elif arg.star == "*":
2534 # Invalid, cannot have "*" follow "**"
2535 raise CSTValidationError(
2536 "Cannot have iterable argument unpacking after keyword argument unpacking."
2537 )
2538 else:
2539 # Invalid, cannot have positional argument follow **/keyword
2540 raise CSTValidationError(
2541 "Cannot have positional argument after keyword argument unpacking."
2542 )
2543
2544 def _check_starred_or_keywords(
2545 self, arg: Arg
2546 ) -> Optional[Callable[[Arg], Callable[[Arg], None]]]:
2547 """
2548 Validates that we only have a mix of "*arg" expansion and "keyword=arg".
2549 """
2550
2551 if arg.keyword is not None:
2552 # Valid, keyword argument
2553 return None
2554 elif arg.star == "**":
2555 # Valid, but we now no longer allow "*" args
2556 # pyre-fixme[7]: Expected `Optional[Callable[[Arg], Callable[...,
2557 # Any]]]` but got `Callable[[Arg], Optional[Callable[[Arg], Callable[...,
2558 # Any]]]]`.
2559 return self._check_kwargs_or_keywords
2560 elif arg.star == "*":
2561 # Valid, iterable unpacking
2562 return None
2563 else:
2564 # Invalid, cannot have positional argument follow **/keyword
2565 raise CSTValidationError(
2566 "Cannot have positional argument after keyword argument."
2567 )
2568
2569 def _check_positional(
2570 self, arg: Arg
2571 ) -> Optional[Callable[[Arg], Callable[[Arg], Callable[[Arg], None]]]]:
2572 """
2573 Validates that we only have a mix of positional args and "*arg" expansion.
2574 """
2575
2576 if arg.keyword is not None:
2577 # Valid, but this puts us into starred/keyword state
2578 # pyre-fixme[7]: Expected `Optional[Callable[[Arg], Callable[...,
2579 # Any]]]` but got `Callable[[Arg], Optional[Callable[[Arg], Callable[...,
2580 # Any]]]]`.
2581 return self._check_starred_or_keywords
2582 elif arg.star == "**":
2583 # Valid, but we skip states to kwargs/keywords
2584 # pyre-fixme[7]: Expected `Optional[Callable[[Arg], Callable[...,
2585 # Any]]]` but got `Callable[[Arg], Optional[Callable[[Arg], Callable[...,
2586 # Any]]]]`.
2587 return self._check_kwargs_or_keywords
2588 elif arg.star == "*":
2589 # Valid, iterator expansion
2590 return None
2591 else:
2592 # Valid, allowed to have positional arguments here
2593 return None
2594
2595 # pyre-fixme[30]: Pyre gave up inferring some types - function `_validate` was
2596 # too complex.
2597 def _validate(self) -> None:
2598 # Validate any super-class stuff, whatever it may be.
2599 super()._validate()
2600 # Now, validate the weird intermingling rules for arguments by running
2601 # a small validator state machine. This works by passing each argument
2602 # to a validator function which can either raise an exception if it
2603 # detects an invalid sequence, return a new validator to be used for the
2604 # next arg, or return None to use the same validator. We could enforce
2605 # always returning ourselves instead of None but it ends up making the
2606 # functions themselves less readable. In this way, the current validator
2607 # function encodes the state we're in (positional state, iterable
2608 # expansion state, or dictionary expansion state).
2609 validator = self._check_positional
2610 for arg in self.args:
2611 validator = validator(arg) or validator
2612
2613
2614@add_slots
2615@dataclass(frozen=True)
2616class Call(_BaseExpressionWithArgs):
2617 """
2618 An expression representing a function call, such as ``do_math(1, 2)`` or
2619 ``picture.post_on_instagram()``.
2620
2621 Function calls consist of a function name and a sequence of arguments wrapped in
2622 :class:`Arg` nodes.
2623 """
2624
2625 #: The expression resulting in a callable that we are to call. Often a :class:`Name`
2626 #: or :class:`Attribute`.
2627 func: BaseExpression
2628
2629 #: The arguments to pass to the resulting callable. These may be a mix of
2630 #: positional arguments, keyword arguments, or "starred" arguments.
2631 args: Sequence[Arg] = ()
2632
2633 lpar: Sequence[LeftParen] = ()
2634 #: Sequence of parenthesis for precedence dictation. These are not the parenthesis
2635 #: before and after the list of ``args``, but rather arguments around the entire
2636 #: call expression, such as ``(( do_math(1, 2) ))``.
2637 rpar: Sequence[RightParen] = ()
2638
2639 #: Whitespace after the ``func`` name, but before the opening parenthesis.
2640 whitespace_after_func: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
2641 #: Whitespace after the opening parenthesis but before the first argument (if there
2642 #: are any). Whitespace after the last argument but before the closing parenthesis
2643 #: is owned by the last :class:`Arg` if it exists.
2644 whitespace_before_args: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
2645
2646 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
2647 """
2648 Calls have a close paren on the right side regardless of whether they're
2649 parenthesized as a whole. As a result, they are safe to use directly against
2650 an adjacent node to the right.
2651 """
2652 if position == ExpressionPosition.LEFT:
2653 return True
2654 if super(Call, self)._safe_to_use_with_word_operator(position):
2655 return True
2656 if position == ExpressionPosition.RIGHT:
2657 return self.func._safe_to_use_with_word_operator(ExpressionPosition.RIGHT)
2658 return False
2659
2660 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Call":
2661 return Call(
2662 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
2663 func=visit_required(self, "func", self.func, visitor),
2664 whitespace_after_func=visit_required(
2665 self, "whitespace_after_func", self.whitespace_after_func, visitor
2666 ),
2667 whitespace_before_args=visit_required(
2668 self, "whitespace_before_args", self.whitespace_before_args, visitor
2669 ),
2670 args=visit_sequence(self, "args", self.args, visitor),
2671 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
2672 )
2673
2674 def _codegen_impl(self, state: CodegenState) -> None:
2675 with self._parenthesize(state):
2676 self.func._codegen(state)
2677 self.whitespace_after_func._codegen(state)
2678 state.add_token("(")
2679 self.whitespace_before_args._codegen(state)
2680 lastarg = len(self.args) - 1
2681 for i, arg in enumerate(self.args):
2682 arg._codegen(state, default_comma=(i != lastarg))
2683 state.add_token(")")
2684
2685
2686@add_slots
2687@dataclass(frozen=True)
2688class Await(BaseExpression):
2689 """
2690 An await expression. Await expressions are only valid inside the body of an
2691 asynchronous :class:`FunctionDef` or (as of Python 3.7) inside of an asynchronous
2692 :class:`GeneratorExp` nodes.
2693 """
2694
2695 #: The actual expression we need to wait for.
2696 expression: BaseExpression
2697
2698 lpar: Sequence[LeftParen] = ()
2699 #: Sequence of parenthesis for precedence dictation.
2700 rpar: Sequence[RightParen] = ()
2701
2702 #: Whitespace that appears after the ``async`` keyword, but before the inner
2703 #: ``expression``.
2704 whitespace_after_await: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
2705
2706 def _validate(self) -> None:
2707 # Validate any super-class stuff, whatever it may be.
2708 super(Await, self)._validate()
2709 # Make sure we don't run identifiers together.
2710 if (
2711 self.whitespace_after_await.empty
2712 and not self.expression._safe_to_use_with_word_operator(
2713 ExpressionPosition.RIGHT
2714 )
2715 ):
2716 raise CSTValidationError("Must have at least one space after await")
2717
2718 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Await":
2719 return Await(
2720 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
2721 whitespace_after_await=visit_required(
2722 self, "whitespace_after_await", self.whitespace_after_await, visitor
2723 ),
2724 expression=visit_required(self, "expression", self.expression, visitor),
2725 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
2726 )
2727
2728 def _codegen_impl(self, state: CodegenState) -> None:
2729 with self._parenthesize(state):
2730 state.add_token("await")
2731 self.whitespace_after_await._codegen(state)
2732 self.expression._codegen(state)
2733
2734
2735@add_slots
2736@dataclass(frozen=True)
2737class IfExp(BaseExpression):
2738 """
2739 An if expression of the form ``body if test else orelse``.
2740
2741 If statements are provided by :class:`If` and :class:`Else` nodes.
2742 """
2743
2744 #: The test to perform.
2745 test: BaseExpression
2746
2747 #: The expression to evaluate when the test is true.
2748 body: BaseExpression
2749
2750 #: The expression to evaluate when the test is false.
2751 orelse: BaseExpression
2752
2753 lpar: Sequence[LeftParen] = ()
2754 #: Sequence of parenthesis for precedence dictation.
2755 rpar: Sequence[RightParen] = ()
2756
2757 #: Whitespace after the ``body`` expression, but before the ``if`` keyword.
2758 whitespace_before_if: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
2759
2760 #: Whitespace after the ``if`` keyword, but before the ``test`` clause.
2761 whitespace_after_if: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
2762
2763 #: Whitespace after the ``test`` expression, but before the ``else`` keyword.
2764 whitespace_before_else: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
2765
2766 #: Whitespace after the ``else`` keyword, but before the ``orelse`` expression.
2767 whitespace_after_else: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
2768
2769 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
2770 if position == ExpressionPosition.RIGHT:
2771 return self.body._safe_to_use_with_word_operator(position)
2772 else:
2773 return self.orelse._safe_to_use_with_word_operator(position)
2774
2775 def _validate(self) -> None:
2776 # Paren validation and such
2777 super(IfExp, self)._validate()
2778 # Validate spacing rules
2779 if (
2780 self.whitespace_before_if.empty
2781 and not self.body._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
2782 ):
2783 raise CSTValidationError(
2784 "Must have at least one space before 'if' keyword."
2785 )
2786 if (
2787 self.whitespace_after_if.empty
2788 and not self.test._safe_to_use_with_word_operator(ExpressionPosition.RIGHT)
2789 ):
2790 raise CSTValidationError("Must have at least one space after 'if' keyword.")
2791 if (
2792 self.whitespace_before_else.empty
2793 and not self.test._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
2794 ):
2795 raise CSTValidationError(
2796 "Must have at least one space before 'else' keyword."
2797 )
2798 if (
2799 self.whitespace_after_else.empty
2800 and not self.orelse._safe_to_use_with_word_operator(
2801 ExpressionPosition.RIGHT
2802 )
2803 ):
2804 raise CSTValidationError(
2805 "Must have at least one space after 'else' keyword."
2806 )
2807
2808 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "IfExp":
2809 return IfExp(
2810 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
2811 body=visit_required(self, "body", self.body, visitor),
2812 whitespace_before_if=visit_required(
2813 self, "whitespace_before_if", self.whitespace_before_if, visitor
2814 ),
2815 whitespace_after_if=visit_required(
2816 self, "whitespace_after_if", self.whitespace_after_if, visitor
2817 ),
2818 test=visit_required(self, "test", self.test, visitor),
2819 whitespace_before_else=visit_required(
2820 self, "whitespace_before_else", self.whitespace_before_else, visitor
2821 ),
2822 whitespace_after_else=visit_required(
2823 self, "whitespace_after_else", self.whitespace_after_else, visitor
2824 ),
2825 orelse=visit_required(self, "orelse", self.orelse, visitor),
2826 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
2827 )
2828
2829 def _codegen_impl(self, state: CodegenState) -> None:
2830 with self._parenthesize(state):
2831 self.body._codegen(state)
2832 self.whitespace_before_if._codegen(state)
2833 state.add_token("if")
2834 self.whitespace_after_if._codegen(state)
2835 self.test._codegen(state)
2836 self.whitespace_before_else._codegen(state)
2837 state.add_token("else")
2838 self.whitespace_after_else._codegen(state)
2839 self.orelse._codegen(state)
2840
2841
2842@add_slots
2843@dataclass(frozen=True)
2844class From(CSTNode):
2845 """
2846 A ``from x`` stanza in a :class:`Yield` or :class:`Raise`.
2847 """
2848
2849 #: The expression that we are yielding/raising from.
2850 item: BaseExpression
2851
2852 #: The whitespace at the very start of this node.
2853 whitespace_before_from: Union[BaseParenthesizableWhitespace, MaybeSentinel] = (
2854 MaybeSentinel.DEFAULT
2855 )
2856
2857 #: The whitespace after the ``from`` keyword, but before the ``item``.
2858 whitespace_after_from: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
2859
2860 def _validate(self) -> None:
2861 if (
2862 isinstance(self.whitespace_after_from, BaseParenthesizableWhitespace)
2863 and self.whitespace_after_from.empty
2864 and not self.item._safe_to_use_with_word_operator(ExpressionPosition.RIGHT)
2865 ):
2866 raise CSTValidationError(
2867 "Must have at least one space after 'from' keyword."
2868 )
2869
2870 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "From":
2871 return From(
2872 whitespace_before_from=visit_sentinel(
2873 self, "whitespace_before_from", self.whitespace_before_from, visitor
2874 ),
2875 whitespace_after_from=visit_required(
2876 self, "whitespace_after_from", self.whitespace_after_from, visitor
2877 ),
2878 item=visit_required(self, "item", self.item, visitor),
2879 )
2880
2881 def _codegen_impl(self, state: CodegenState, default_space: str = "") -> None:
2882 whitespace_before_from = self.whitespace_before_from
2883 if isinstance(whitespace_before_from, BaseParenthesizableWhitespace):
2884 whitespace_before_from._codegen(state)
2885 else:
2886 state.add_token(default_space)
2887
2888 with state.record_syntactic_position(self):
2889 state.add_token("from")
2890 self.whitespace_after_from._codegen(state)
2891 self.item._codegen(state)
2892
2893
2894@add_slots
2895@dataclass(frozen=True)
2896class Yield(BaseExpression):
2897 """
2898 A yield expression similar to ``yield x`` or ``yield from fun()``.
2899
2900 To learn more about the ways that yield can be used in generators, refer to
2901 `Python's language reference
2902 <https://docs.python.org/3/reference/expressions.html#yieldexpr>`__.
2903 """
2904
2905 #: The value yielded from the generator, in the case of a :class:`From` clause, a
2906 #: sub-generator to iterate over.
2907 value: Optional[Union[BaseExpression, From]] = None
2908
2909 lpar: Sequence[LeftParen] = ()
2910 #: Sequence of parenthesis for precedence dictation.
2911 rpar: Sequence[RightParen] = ()
2912
2913 #: Whitespace after the ``yield`` keyword, but before the ``value``.
2914 whitespace_after_yield: Union[BaseParenthesizableWhitespace, MaybeSentinel] = (
2915 MaybeSentinel.DEFAULT
2916 )
2917
2918 def _validate(self) -> None:
2919 # Paren rules and such
2920 super(Yield, self)._validate()
2921 # Our own rules
2922 whitespace_after_yield = self.whitespace_after_yield
2923 if (
2924 isinstance(whitespace_after_yield, BaseParenthesizableWhitespace)
2925 and whitespace_after_yield.empty
2926 ):
2927 value = self.value
2928 if isinstance(value, From):
2929 raise CSTValidationError(
2930 "Must have at least one space after 'yield' keyword."
2931 )
2932 if isinstance(
2933 value, BaseExpression
2934 ) and not value._safe_to_use_with_word_operator(ExpressionPosition.RIGHT):
2935 raise CSTValidationError(
2936 "Must have at least one space after 'yield' keyword."
2937 )
2938
2939 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Yield":
2940 return Yield(
2941 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
2942 whitespace_after_yield=visit_sentinel(
2943 self, "whitespace_after_yield", self.whitespace_after_yield, visitor
2944 ),
2945 value=visit_optional(self, "value", self.value, visitor),
2946 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
2947 )
2948
2949 def _codegen_impl(self, state: CodegenState) -> None:
2950 with self._parenthesize(state):
2951 state.add_token("yield")
2952 whitespace_after_yield = self.whitespace_after_yield
2953 if isinstance(whitespace_after_yield, BaseParenthesizableWhitespace):
2954 whitespace_after_yield._codegen(state)
2955 else:
2956 # Only need a space after yield if there is a value to yield.
2957 if self.value is not None:
2958 state.add_token(" ")
2959 value = self.value
2960 if isinstance(value, From):
2961 value._codegen(state, default_space="")
2962 elif value is not None:
2963 value._codegen(state)
2964
2965
2966class _BaseElementImpl(CSTNode, ABC):
2967 """
2968 An internal base class for :class:`Element` and :class:`DictElement`.
2969 """
2970
2971 __slots__ = ()
2972
2973 value: BaseExpression
2974 comma: Union[Comma, MaybeSentinel] = MaybeSentinel.DEFAULT
2975
2976 def _codegen_comma(
2977 self,
2978 state: CodegenState,
2979 default_comma: bool = False,
2980 default_comma_whitespace: bool = False, # False for a single-item collection
2981 ) -> None:
2982 """
2983 Called by `_codegen_impl` in subclasses to generate the comma.
2984 """
2985 comma = self.comma
2986 if comma is MaybeSentinel.DEFAULT and default_comma:
2987 if default_comma_whitespace:
2988 state.add_token(", ")
2989 else:
2990 state.add_token(",")
2991 elif isinstance(comma, Comma):
2992 comma._codegen(state)
2993
2994 @abstractmethod
2995 def _codegen_impl(
2996 self,
2997 state: CodegenState,
2998 default_comma: bool = False,
2999 default_comma_whitespace: bool = False, # False for a single-item collection
3000 ) -> None: ...
3001
3002
3003class BaseElement(_BaseElementImpl, ABC):
3004 """
3005 An element of a literal list, tuple, or set. For elements of a literal dict, see
3006 BaseDictElement.
3007 """
3008
3009 __slots__ = ()
3010
3011
3012class BaseDictElement(_BaseElementImpl, ABC):
3013 """
3014 An element of a literal dict. For elements of a list, tuple, or set, see
3015 BaseElement.
3016 """
3017
3018 __slots__ = ()
3019
3020
3021@add_slots
3022@dataclass(frozen=True)
3023class Element(BaseElement):
3024 """
3025 A simple value in a literal :class:`List`, :class:`Tuple`, or :class:`Set`.
3026 These a literal collection may also contain a :class:`StarredElement`.
3027
3028 If you're using a literal :class:`Dict`, see :class:`DictElement` instead.
3029 """
3030
3031 value: BaseExpression
3032
3033 #: A trailing comma. By default, we'll only insert a comma if one is required.
3034 comma: Union[Comma, MaybeSentinel] = MaybeSentinel.DEFAULT
3035
3036 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Element":
3037 return Element(
3038 value=visit_required(self, "value", self.value, visitor),
3039 comma=visit_sentinel(self, "comma", self.comma, visitor),
3040 )
3041
3042 def _codegen_impl(
3043 self,
3044 state: CodegenState,
3045 default_comma: bool = False,
3046 default_comma_whitespace: bool = False,
3047 ) -> None:
3048 with state.record_syntactic_position(self):
3049 self.value._codegen(state)
3050 self._codegen_comma(state, default_comma, default_comma_whitespace)
3051
3052
3053@add_slots
3054@dataclass(frozen=True)
3055class DictElement(BaseDictElement):
3056 """
3057 A simple ``key: value`` pair that represents a single entry in a literal
3058 :class:`Dict`. :class:`Dict` nodes may also contain a
3059 :class:`StarredDictElement`.
3060
3061 If you're using a literal :class:`List`, :class:`Tuple`, or :class:`Set`,
3062 see :class:`Element` instead.
3063 """
3064
3065 key: BaseExpression
3066 value: BaseExpression
3067
3068 #: A trailing comma. By default, we'll only insert a comma if one is required.
3069 comma: Union[Comma, MaybeSentinel] = MaybeSentinel.DEFAULT
3070
3071 #: Whitespace after the key, but before the colon in ``key : value``.
3072 whitespace_before_colon: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
3073 #: Whitespace after the colon, but before the value in ``key : value``.
3074 whitespace_after_colon: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
3075
3076 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "DictElement":
3077 return DictElement(
3078 key=visit_required(self, "key", self.key, visitor),
3079 whitespace_before_colon=visit_required(
3080 self, "whitespace_before_colon", self.whitespace_before_colon, visitor
3081 ),
3082 whitespace_after_colon=visit_required(
3083 self, "whitespace_after_colon", self.whitespace_after_colon, visitor
3084 ),
3085 value=visit_required(self, "value", self.value, visitor),
3086 comma=visit_sentinel(self, "comma", self.comma, visitor),
3087 )
3088
3089 def _codegen_impl(
3090 self,
3091 state: CodegenState,
3092 default_comma: bool = False,
3093 default_comma_whitespace: bool = False,
3094 ) -> None:
3095 with state.record_syntactic_position(self):
3096 self.key._codegen(state)
3097 self.whitespace_before_colon._codegen(state)
3098 state.add_token(":")
3099 self.whitespace_after_colon._codegen(state)
3100 self.value._codegen(state)
3101 self._codegen_comma(state, default_comma, default_comma_whitespace)
3102
3103
3104@add_slots
3105@dataclass(frozen=True)
3106class StarredElement(BaseElement, BaseExpression, _BaseParenthesizedNode):
3107 """
3108 A starred ``*value`` element that expands to represent multiple values in a literal
3109 :class:`List`, :class:`Tuple`, or :class:`Set`.
3110
3111 If you're using a literal :class:`Dict`, see :class:`StarredDictElement` instead.
3112
3113 If this node owns parenthesis, those parenthesis wrap the leading asterisk, but not
3114 the trailing comma. For example::
3115
3116 StarredElement(
3117 cst.Name("el"),
3118 comma=cst.Comma(),
3119 lpar=[cst.LeftParen()],
3120 rpar=[cst.RightParen()],
3121 )
3122
3123 will generate::
3124
3125 (*el),
3126 """
3127
3128 value: BaseExpression
3129
3130 #: A trailing comma. By default, we'll only insert a comma if one is required.
3131 comma: Union[Comma, MaybeSentinel] = MaybeSentinel.DEFAULT
3132
3133 #: Parenthesis at the beginning of the node, before the leading asterisk.
3134 lpar: Sequence[LeftParen] = ()
3135 #: Parentheses after the value, but before a comma (if there is one).
3136 rpar: Sequence[RightParen] = ()
3137
3138 #: Whitespace between the leading asterisk and the value expression.
3139 whitespace_before_value: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
3140
3141 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "StarredElement":
3142 return StarredElement(
3143 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
3144 whitespace_before_value=visit_required(
3145 self, "whitespace_before_value", self.whitespace_before_value, visitor
3146 ),
3147 value=visit_required(self, "value", self.value, visitor),
3148 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
3149 comma=visit_sentinel(self, "comma", self.comma, visitor),
3150 )
3151
3152 def _codegen_impl(
3153 self,
3154 state: CodegenState,
3155 default_comma: bool = False,
3156 default_comma_whitespace: bool = False,
3157 ) -> None:
3158 with self._parenthesize(state):
3159 state.add_token("*")
3160 self.whitespace_before_value._codegen(state)
3161 self.value._codegen(state)
3162 self._codegen_comma(state, default_comma, default_comma_whitespace)
3163
3164
3165@add_slots
3166@dataclass(frozen=True)
3167class StarredDictElement(BaseDictElement):
3168 """
3169 A starred ``**value`` element that expands to represent multiple values in a literal
3170 :class:`Dict`.
3171
3172 If you're using a literal :class:`List`, :class:`Tuple`, or :class:`Set`,
3173 see :class:`StarredElement` instead.
3174
3175 Unlike :class:`StarredElement`, this node does not own left or right parenthesis,
3176 but the ``value`` field may still contain parenthesis. This is due to some
3177 asymmetry in Python's grammar.
3178 """
3179
3180 value: BaseExpression
3181
3182 #: A trailing comma. By default, we'll only insert a comma if one is required.
3183 comma: Union[Comma, MaybeSentinel] = MaybeSentinel.DEFAULT
3184
3185 #: Whitespace between the leading asterisks and the value expression.
3186 whitespace_before_value: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
3187
3188 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "StarredDictElement":
3189 return StarredDictElement(
3190 whitespace_before_value=visit_required(
3191 self, "whitespace_before_value", self.whitespace_before_value, visitor
3192 ),
3193 value=visit_required(self, "value", self.value, visitor),
3194 comma=visit_sentinel(self, "comma", self.comma, visitor),
3195 )
3196
3197 def _codegen_impl(
3198 self,
3199 state: CodegenState,
3200 default_comma: bool = False,
3201 default_comma_whitespace: bool = False,
3202 ) -> None:
3203 with state.record_syntactic_position(self):
3204 state.add_token("**")
3205 self.whitespace_before_value._codegen(state)
3206 self.value._codegen(state)
3207 self._codegen_comma(state, default_comma, default_comma_whitespace)
3208
3209
3210@add_slots
3211@dataclass(frozen=True)
3212class Tuple(BaseAssignTargetExpression, BaseDelTargetExpression):
3213 """
3214 An immutable literal tuple. Tuples are often (but not always) parenthesized.
3215
3216 ::
3217
3218 Tuple([
3219 Element(Integer("1")),
3220 Element(Integer("2")),
3221 StarredElement(Name("others")),
3222 ])
3223
3224 generates the following code::
3225
3226 (1, 2, *others)
3227 """
3228
3229 #: A sequence containing all the :class:`Element` and :class:`StarredElement` nodes
3230 #: in the tuple.
3231 elements: Sequence[BaseElement]
3232
3233 lpar: Sequence[LeftParen] = field(default_factory=lambda: (LeftParen(),))
3234 #: Sequence of parenthesis for precedence dictation.
3235 rpar: Sequence[RightParen] = field(default_factory=lambda: (RightParen(),))
3236
3237 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
3238 if super(Tuple, self)._safe_to_use_with_word_operator(position):
3239 # if we have parenthesis, we're safe.
3240 return True
3241 # elements[-1] and elements[0] must exist past this point, because
3242 # we're not parenthesized, meaning we must have at least one element.
3243 elements = self.elements
3244 if position == ExpressionPosition.LEFT:
3245 last_element = elements[-1]
3246 return (
3247 isinstance(last_element.comma, Comma)
3248 or (
3249 isinstance(last_element, StarredElement)
3250 and len(last_element.rpar) > 0
3251 )
3252 or last_element.value._safe_to_use_with_word_operator(position)
3253 )
3254 else: # ExpressionPosition.RIGHT
3255 first_element = elements[0]
3256 # starred elements are always safe because they begin with ( or *
3257 return isinstance(
3258 first_element, StarredElement
3259 ) or first_element.value._safe_to_use_with_word_operator(position)
3260
3261 def _validate(self) -> None:
3262 # Paren validation and such
3263 super(Tuple, self)._validate()
3264
3265 if len(self.elements) == 0:
3266 if len(self.lpar) == 0: # assumes len(lpar) == len(rpar), via superclass
3267 raise CSTValidationError(
3268 "A zero-length tuple must be wrapped in parentheses."
3269 )
3270 # Invalid commas aren't possible, because MaybeSentinel will ensure that there
3271 # is a comma where required.
3272
3273 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Tuple":
3274 return Tuple(
3275 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
3276 elements=visit_sequence(self, "elements", self.elements, visitor),
3277 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
3278 )
3279
3280 def _codegen_impl(self, state: CodegenState) -> None:
3281 with self._parenthesize(state):
3282 elements = self.elements
3283 if len(elements) == 1:
3284 elements[0]._codegen(
3285 state, default_comma=True, default_comma_whitespace=False
3286 )
3287 else:
3288 for idx, el in enumerate(elements):
3289 el._codegen(
3290 state,
3291 default_comma=(idx < len(elements) - 1),
3292 default_comma_whitespace=True,
3293 )
3294
3295
3296class BaseList(BaseExpression, ABC):
3297 """
3298 A base class for :class:`List` and :class:`ListComp`, which both result in a list
3299 object when evaluated.
3300 """
3301
3302 __slots__ = ()
3303
3304 lbracket: LeftSquareBracket = LeftSquareBracket.field()
3305 #: Brackets surrounding the list.
3306 rbracket: RightSquareBracket = RightSquareBracket.field()
3307
3308 lpar: Sequence[LeftParen] = ()
3309 #: Sequence of parenthesis for precedence dictation.
3310 rpar: Sequence[RightParen] = ()
3311
3312 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
3313 return True
3314
3315 @contextmanager
3316 def _bracketize(self, state: CodegenState) -> Generator[None, None, None]:
3317 self.lbracket._codegen(state)
3318 yield
3319 self.rbracket._codegen(state)
3320
3321
3322@add_slots
3323@dataclass(frozen=True)
3324class List(BaseList, BaseAssignTargetExpression, BaseDelTargetExpression):
3325 """
3326 A mutable literal list.
3327
3328 ::
3329
3330 List([
3331 Element(Integer("1")),
3332 Element(Integer("2")),
3333 StarredElement(Name("others")),
3334 ])
3335
3336 generates the following code::
3337
3338 [1, 2, *others]
3339
3340 List comprehensions are represented with a :class:`ListComp` node.
3341 """
3342
3343 #: A sequence containing all the :class:`Element` and :class:`StarredElement` nodes
3344 #: in the list.
3345 elements: Sequence[BaseElement]
3346
3347 lbracket: LeftSquareBracket = LeftSquareBracket.field()
3348 #: Brackets surrounding the list.
3349 rbracket: RightSquareBracket = RightSquareBracket.field()
3350
3351 lpar: Sequence[LeftParen] = ()
3352 #: Sequence of parenthesis for precedence dictation.
3353 rpar: Sequence[RightParen] = ()
3354
3355 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "List":
3356 return List(
3357 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
3358 lbracket=visit_required(self, "lbracket", self.lbracket, visitor),
3359 elements=visit_sequence(self, "elements", self.elements, visitor),
3360 rbracket=visit_required(self, "rbracket", self.rbracket, visitor),
3361 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
3362 )
3363
3364 def _codegen_impl(self, state: CodegenState) -> None:
3365 with self._parenthesize(state), self._bracketize(state):
3366 elements = self.elements
3367 for idx, el in enumerate(elements):
3368 el._codegen(
3369 state,
3370 default_comma=(idx < len(elements) - 1),
3371 default_comma_whitespace=True,
3372 )
3373
3374
3375class _BaseSetOrDict(BaseExpression, ABC):
3376 """
3377 An abstract base class for :class:`BaseSet` and :class:`BaseDict`.
3378
3379 Literal sets and dicts are syntactically similar (hence this shared base class), but
3380 are semantically different. This base class is an implementation detail and
3381 shouldn't be exported.
3382 """
3383
3384 __slots__ = ()
3385
3386 lbrace: LeftCurlyBrace = LeftCurlyBrace.field()
3387 #: Braces surrounding the set or dict.
3388 rbrace: RightCurlyBrace = RightCurlyBrace.field()
3389
3390 lpar: Sequence[LeftParen] = ()
3391 #: Sequence of parenthesis for precedence dictation.
3392 rpar: Sequence[RightParen] = ()
3393
3394 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
3395 return True
3396
3397 # brace-ize seems like a very made-up word. And it is!
3398 @contextmanager
3399 def _braceize(self, state: CodegenState) -> Generator[None, None, None]:
3400 self.lbrace._codegen(state)
3401 yield
3402 self.rbrace._codegen(state)
3403
3404
3405class BaseSet(_BaseSetOrDict, ABC):
3406 """
3407 An abstract base class for :class:`Set` and :class:`SetComp`, which both result in
3408 a set object when evaluated.
3409 """
3410
3411 __slots__ = ()
3412
3413
3414@add_slots
3415@dataclass(frozen=True)
3416class Set(BaseSet):
3417 """
3418 A mutable literal set.
3419
3420 ::
3421
3422 Set([
3423 Element(Integer("1")),
3424 Element(Integer("2")),
3425 StarredElement(Name("others")),
3426 ])
3427
3428 generates the following code::
3429
3430 {1, 2, *others}
3431
3432 Set comprehensions are represented with a :class:`SetComp` node.
3433 """
3434
3435 #: A sequence containing all the :class:`Element` and :class:`StarredElement` nodes
3436 #: in the set.
3437 elements: Sequence[BaseElement]
3438
3439 lbrace: LeftCurlyBrace = LeftCurlyBrace.field()
3440 #: Braces surrounding the set.
3441 rbrace: RightCurlyBrace = RightCurlyBrace.field()
3442
3443 lpar: Sequence[LeftParen] = ()
3444 #: Sequence of parenthesis for precedence dictation.
3445 rpar: Sequence[RightParen] = ()
3446
3447 def _validate(self) -> None:
3448 super(Set, self)._validate()
3449
3450 if len(self.elements) == 0:
3451 raise CSTValidationError(
3452 "A literal set must have at least one element. A zero-element set "
3453 + "would be syntatically ambiguous with an empty dict, `{}`."
3454 )
3455
3456 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Set":
3457 return Set(
3458 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
3459 lbrace=visit_required(self, "lbrace", self.lbrace, visitor),
3460 elements=visit_sequence(self, "elements", self.elements, visitor),
3461 rbrace=visit_required(self, "rbrace", self.rbrace, visitor),
3462 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
3463 )
3464
3465 def _codegen_impl(self, state: CodegenState) -> None:
3466 with self._parenthesize(state), self._braceize(state):
3467 elements = self.elements
3468 for idx, el in enumerate(elements):
3469 el._codegen(
3470 state,
3471 default_comma=(idx < len(elements) - 1),
3472 default_comma_whitespace=True,
3473 )
3474
3475
3476class BaseDict(_BaseSetOrDict, ABC):
3477 """
3478 An abstract base class for :class:`Dict` and :class:`DictComp`, which both result in
3479 a dict object when evaluated.
3480 """
3481
3482 __slots__ = ()
3483
3484
3485@add_slots
3486@dataclass(frozen=True)
3487class Dict(BaseDict):
3488 """
3489 A literal dictionary. Key-value pairs are stored in ``elements`` using
3490 :class:`DictElement` nodes.
3491
3492 It's possible to expand one dictionary into another, as in ``{k: v, **expanded}``.
3493 Expanded elements are stored as :class:`StarredDictElement` nodes.
3494
3495 ::
3496
3497 Dict([
3498 DictElement(Name("k1"), Name("v1")),
3499 DictElement(Name("k2"), Name("v2")),
3500 StarredDictElement(Name("expanded")),
3501 ])
3502
3503 generates the following code::
3504
3505 {k1: v1, k2: v2, **expanded}
3506 """
3507
3508 elements: Sequence[BaseDictElement]
3509 lbrace: LeftCurlyBrace = LeftCurlyBrace.field()
3510 rbrace: RightCurlyBrace = RightCurlyBrace.field()
3511 lpar: Sequence[LeftParen] = ()
3512 rpar: Sequence[RightParen] = ()
3513
3514 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "Dict":
3515 return Dict(
3516 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
3517 lbrace=visit_required(self, "lbrace", self.lbrace, visitor),
3518 elements=visit_sequence(self, "elements", self.elements, visitor),
3519 rbrace=visit_required(self, "rbrace", self.rbrace, visitor),
3520 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
3521 )
3522
3523 def _codegen_impl(self, state: CodegenState) -> None:
3524 with self._parenthesize(state), self._braceize(state):
3525 elements = self.elements
3526 for idx, el in enumerate(elements):
3527 el._codegen(
3528 state,
3529 default_comma=(idx < len(elements) - 1),
3530 default_comma_whitespace=True,
3531 )
3532
3533
3534@add_slots
3535@dataclass(frozen=True)
3536class CompFor(CSTNode):
3537 """
3538 One ``for`` clause in a :class:`BaseComp`, or a nested hierarchy of
3539 ``for`` clauses.
3540
3541 Nested loops in comprehensions are difficult to get right, but they can be thought
3542 of as a flat representation of nested clauses.
3543
3544 ``elt for a in b for c in d if e`` can be thought of as::
3545
3546 for a in b:
3547 for c in d:
3548 if e:
3549 yield elt
3550
3551 And that would form the following CST::
3552
3553 ListComp(
3554 elt=Name("elt"),
3555 for_in=CompFor(
3556 target=Name("a"),
3557 iter=Name("b"),
3558 ifs=[],
3559 inner_comp_for=CompFor(
3560 target=Name("c"),
3561 iter=Name("d"),
3562 ifs=[
3563 CompIf(
3564 test=Name("e"),
3565 ),
3566 ],
3567 ),
3568 ),
3569 )
3570
3571 Normal ``for`` statements are provided by :class:`For`.
3572 """
3573
3574 #: The target to assign a value to in each iteration of the loop. This is different
3575 #: from :attr:`GeneratorExp.elt`, :attr:`ListComp.elt`, :attr:`SetComp.elt`, and
3576 #: ``key`` and ``value`` in :class:`DictComp`, because it doesn't directly effect
3577 #: the value of resulting generator, list, set, or dict.
3578 target: BaseAssignTargetExpression
3579
3580 #: The value to iterate over. Every value in ``iter`` is stored in ``target``.
3581 iter: BaseExpression
3582
3583 #: Zero or more conditional clauses that control this loop. If any of these tests
3584 #: fail, the ``target`` item is skipped.
3585 #:
3586 #: ::
3587 #:
3588 #: if a if b if c
3589 #:
3590 #: has similar semantics to::
3591 #:
3592 #: if a and b and c
3593 ifs: Sequence["CompIf"] = ()
3594
3595 #: Another :class:`CompFor` node used to form nested loops. Nested comprehensions
3596 #: can be useful, but they tend to be difficult to read and write. As a result they
3597 #: are uncommon.
3598 inner_for_in: Optional["CompFor"] = None
3599
3600 #: An optional async modifier that appears before the ``for`` keyword.
3601 asynchronous: Optional[Asynchronous] = None
3602
3603 #: Whitespace that appears at the beginning of this node, before the ``for`` and
3604 #: ``async`` keywords.
3605 whitespace_before: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
3606
3607 #: Whitespace appearing after the ``for`` keyword, but before the ``target``.
3608 whitespace_after_for: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
3609
3610 #: Whitespace appearing after the ``target``, but before the ``in`` keyword.
3611 whitespace_before_in: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
3612
3613 #: Whitespace appearing after the ``in`` keyword, but before the ``iter``.
3614 whitespace_after_in: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
3615
3616 def _validate(self) -> None:
3617 if (
3618 self.whitespace_after_for.empty
3619 and not self.target._safe_to_use_with_word_operator(
3620 ExpressionPosition.RIGHT
3621 )
3622 ):
3623 raise CSTValidationError(
3624 "Must have at least one space after 'for' keyword."
3625 )
3626
3627 if (
3628 self.whitespace_before_in.empty
3629 and not self.target._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
3630 ):
3631 raise CSTValidationError(
3632 "Must have at least one space before 'in' keyword."
3633 )
3634
3635 if (
3636 self.whitespace_after_in.empty
3637 and not self.iter._safe_to_use_with_word_operator(ExpressionPosition.RIGHT)
3638 ):
3639 raise CSTValidationError("Must have at least one space after 'in' keyword.")
3640
3641 prev_expr = self.iter
3642 for if_clause in self.ifs:
3643 if (
3644 if_clause.whitespace_before.empty
3645 and not prev_expr._safe_to_use_with_word_operator(
3646 ExpressionPosition.LEFT
3647 )
3648 ):
3649 raise CSTValidationError(
3650 "Must have at least one space before 'if' keyword."
3651 )
3652 prev_expr = if_clause.test
3653
3654 inner_for_in = self.inner_for_in
3655 if (
3656 inner_for_in is not None
3657 and inner_for_in.whitespace_before.empty
3658 and not prev_expr._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
3659 ):
3660 keyword = "async" if inner_for_in.asynchronous else "for"
3661 raise CSTValidationError(
3662 f"Must have at least one space before '{keyword}' keyword."
3663 )
3664
3665 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "CompFor":
3666 return CompFor(
3667 whitespace_before=visit_required(
3668 self, "whitespace_before", self.whitespace_before, visitor
3669 ),
3670 asynchronous=visit_optional(
3671 self, "asynchronous", self.asynchronous, visitor
3672 ),
3673 whitespace_after_for=visit_required(
3674 self, "whitespace_after_for", self.whitespace_after_for, visitor
3675 ),
3676 target=visit_required(self, "target", self.target, visitor),
3677 whitespace_before_in=visit_required(
3678 self, "whitespace_before_in", self.whitespace_before_in, visitor
3679 ),
3680 whitespace_after_in=visit_required(
3681 self, "whitespace_after_in", self.whitespace_after_in, visitor
3682 ),
3683 iter=visit_required(self, "iter", self.iter, visitor),
3684 ifs=visit_sequence(self, "ifs", self.ifs, visitor),
3685 inner_for_in=visit_optional(
3686 self, "inner_for_in", self.inner_for_in, visitor
3687 ),
3688 )
3689
3690 def _codegen_impl(self, state: CodegenState) -> None:
3691 self.whitespace_before._codegen(state)
3692 asynchronous = self.asynchronous
3693 if asynchronous is not None:
3694 asynchronous._codegen(state)
3695 state.add_token("for")
3696 self.whitespace_after_for._codegen(state)
3697 self.target._codegen(state)
3698 self.whitespace_before_in._codegen(state)
3699 state.add_token("in")
3700 self.whitespace_after_in._codegen(state)
3701 self.iter._codegen(state)
3702 ifs = self.ifs
3703 for if_clause in ifs:
3704 if_clause._codegen(state)
3705 inner_for_in = self.inner_for_in
3706 if inner_for_in is not None:
3707 inner_for_in._codegen(state)
3708
3709
3710@add_slots
3711@dataclass(frozen=True)
3712class CompIf(CSTNode):
3713 """
3714 A conditional clause in a :class:`CompFor`, used as part of a generator or
3715 comprehension expression.
3716
3717 If the ``test`` fails, the current element in the :class:`CompFor` will be skipped.
3718 """
3719
3720 #: An expression to evaluate. When interpreted, Python will coerce it to a boolean.
3721 test: BaseExpression
3722
3723 #: Whitespace before the ``if`` keyword.
3724 whitespace_before: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
3725
3726 #: Whitespace after the ``if`` keyword, but before the ``test`` expression.
3727 whitespace_before_test: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
3728
3729 def _validate(self) -> None:
3730 if (
3731 self.whitespace_before_test.empty
3732 and not self.test._safe_to_use_with_word_operator(ExpressionPosition.RIGHT)
3733 ):
3734 raise CSTValidationError("Must have at least one space after 'if' keyword.")
3735
3736 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "CompIf":
3737 return CompIf(
3738 whitespace_before=visit_required(
3739 self, "whitespace_before", self.whitespace_before, visitor
3740 ),
3741 whitespace_before_test=visit_required(
3742 self, "whitespace_before_test", self.whitespace_before_test, visitor
3743 ),
3744 test=visit_required(self, "test", self.test, visitor),
3745 )
3746
3747 def _codegen_impl(self, state: CodegenState) -> None:
3748 self.whitespace_before._codegen(state)
3749 state.add_token("if")
3750 self.whitespace_before_test._codegen(state)
3751 self.test._codegen(state)
3752
3753
3754class BaseComp(BaseExpression, ABC):
3755 """
3756 A base class for all comprehension and generator expressions, including
3757 :class:`GeneratorExp`, :class:`ListComp`, :class:`SetComp`, and :class:`DictComp`.
3758 """
3759
3760 __slots__ = ()
3761
3762 for_in: CompFor
3763
3764
3765class BaseSimpleComp(BaseComp, ABC):
3766 """
3767 The base class for :class:`ListComp`, :class:`SetComp`, and :class:`GeneratorExp`.
3768 :class:`DictComp` is not a :class:`BaseSimpleComp`, because it uses ``key`` and
3769 ``value``.
3770 """
3771
3772 __slots__ = ()
3773
3774 #: The expression evaluated during each iteration of the comprehension. This
3775 #: lexically comes before the ``for_in`` clause, but it is semantically the
3776 #: inner-most element, evaluated inside the ``for_in`` clause.
3777 elt: BaseExpression
3778
3779 #: The ``for ... in ... if ...`` clause that lexically comes after ``elt``. This may
3780 #: be a nested structure for nested comprehensions. See :class:`CompFor` for
3781 #: details.
3782 for_in: CompFor
3783
3784 def _validate(self) -> None:
3785 super(BaseSimpleComp, self)._validate()
3786
3787 for_in = self.for_in
3788 if (
3789 for_in.whitespace_before.empty
3790 and not self.elt._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
3791 ):
3792 keyword = "async" if for_in.asynchronous else "for"
3793 raise CSTValidationError(
3794 f"Must have at least one space before '{keyword}' keyword."
3795 )
3796
3797
3798@add_slots
3799@dataclass(frozen=True)
3800class GeneratorExp(BaseSimpleComp):
3801 """
3802 A generator expression. ``elt`` represents the value yielded for each item in
3803 :attr:`CompFor.iter`.
3804
3805 All ``for ... in ...`` and ``if ...`` clauses are stored as a recursive
3806 :class:`CompFor` data structure inside ``for_in``.
3807 """
3808
3809 #: The expression evaluated and yielded during each iteration of the generator.
3810 elt: BaseExpression
3811
3812 #: The ``for ... in ... if ...`` clause that comes after ``elt``. This may be a
3813 #: nested structure for nested comprehensions. See :class:`CompFor` for details.
3814 for_in: CompFor
3815
3816 lpar: Sequence[LeftParen] = field(default_factory=lambda: (LeftParen(),))
3817 #: Sequence of parentheses for precedence dictation. Generator expressions must
3818 #: always be parenthesized. However, if a generator expression is the only argument
3819 #: inside a function call, the enclosing :class:`Call` node may own the parentheses
3820 #: instead.
3821 rpar: Sequence[RightParen] = field(default_factory=lambda: (RightParen(),))
3822
3823 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
3824 # Generators are always parenthesized
3825 return True
3826
3827 # A note about validation: Generators must always be parenthesized, but it's
3828 # possible that this Generator node doesn't own those parenthesis (in the case of a
3829 # function call with a single generator argument).
3830 #
3831 # Therefore, there's no useful validation we can do here. In theory, our parent
3832 # could do the validation, but there's a ton of potential parents to a Generator, so
3833 # it's not worth the effort.
3834
3835 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "GeneratorExp":
3836 return GeneratorExp(
3837 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
3838 elt=visit_required(self, "elt", self.elt, visitor),
3839 for_in=visit_required(self, "for_in", self.for_in, visitor),
3840 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
3841 )
3842
3843 def _codegen_impl(self, state: CodegenState) -> None:
3844 with self._parenthesize(state):
3845 self.elt._codegen(state)
3846 self.for_in._codegen(state)
3847
3848
3849@add_slots
3850@dataclass(frozen=True)
3851class ListComp(BaseList, BaseSimpleComp):
3852 """
3853 A list comprehension. ``elt`` represents the value stored for each item in
3854 :attr:`CompFor.iter`.
3855
3856 All ``for ... in ...`` and ``if ...`` clauses are stored as a recursive
3857 :class:`CompFor` data structure inside ``for_in``.
3858 """
3859
3860 #: The expression evaluated and stored during each iteration of the comprehension.
3861 elt: BaseExpression
3862
3863 #: The ``for ... in ... if ...`` clause that comes after ``elt``. This may be a
3864 #: nested structure for nested comprehensions. See :class:`CompFor` for details.
3865 for_in: CompFor
3866
3867 lbracket: LeftSquareBracket = LeftSquareBracket.field()
3868 #: Brackets surrounding the list comprehension.
3869 rbracket: RightSquareBracket = RightSquareBracket.field()
3870
3871 lpar: Sequence[LeftParen] = ()
3872 #: Sequence of parenthesis for precedence dictation.
3873 rpar: Sequence[RightParen] = ()
3874
3875 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "ListComp":
3876 return ListComp(
3877 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
3878 lbracket=visit_required(self, "lbracket", self.lbracket, visitor),
3879 elt=visit_required(self, "elt", self.elt, visitor),
3880 for_in=visit_required(self, "for_in", self.for_in, visitor),
3881 rbracket=visit_required(self, "rbracket", self.rbracket, visitor),
3882 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
3883 )
3884
3885 def _codegen_impl(self, state: CodegenState) -> None:
3886 with self._parenthesize(state), self._bracketize(state):
3887 self.elt._codegen(state)
3888 self.for_in._codegen(state)
3889
3890
3891@add_slots
3892@dataclass(frozen=True)
3893class SetComp(BaseSet, BaseSimpleComp):
3894 """
3895 A set comprehension. ``elt`` represents the value stored for each item in
3896 :attr:`CompFor.iter`.
3897
3898 All ``for ... in ...`` and ``if ...`` clauses are stored as a recursive
3899 :class:`CompFor` data structure inside ``for_in``.
3900 """
3901
3902 #: The expression evaluated and stored during each iteration of the comprehension.
3903 elt: BaseExpression
3904
3905 #: The ``for ... in ... if ...`` clause that comes after ``elt``. This may be a
3906 #: nested structure for nested comprehensions. See :class:`CompFor` for details.
3907 for_in: CompFor
3908
3909 lbrace: LeftCurlyBrace = LeftCurlyBrace.field()
3910 #: Braces surrounding the set comprehension.
3911 rbrace: RightCurlyBrace = RightCurlyBrace.field()
3912
3913 lpar: Sequence[LeftParen] = ()
3914 #: Sequence of parenthesis for precedence dictation.
3915 rpar: Sequence[RightParen] = ()
3916
3917 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "SetComp":
3918 return SetComp(
3919 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
3920 lbrace=visit_required(self, "lbrace", self.lbrace, visitor),
3921 elt=visit_required(self, "elt", self.elt, visitor),
3922 for_in=visit_required(self, "for_in", self.for_in, visitor),
3923 rbrace=visit_required(self, "rbrace", self.rbrace, visitor),
3924 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
3925 )
3926
3927 def _codegen_impl(self, state: CodegenState) -> None:
3928 with self._parenthesize(state), self._braceize(state):
3929 self.elt._codegen(state)
3930 self.for_in._codegen(state)
3931
3932
3933@add_slots
3934@dataclass(frozen=True)
3935class DictComp(BaseDict, BaseComp):
3936 """
3937 A dictionary comprehension. ``key`` and ``value`` represent the dictionary entry
3938 evaluated for each item.
3939
3940 All ``for ... in ...`` and ``if ...`` clauses are stored as a recursive
3941 :class:`CompFor` data structure inside ``for_in``.
3942 """
3943
3944 #: The key inserted into the dictionary during each iteration of the comprehension.
3945 key: BaseExpression
3946 #: The value associated with the ``key`` inserted into the dictionary during each
3947 #: iteration of the comprehension.
3948 value: BaseExpression
3949
3950 #: The ``for ... in ... if ...`` clause that lexically comes after ``key`` and
3951 #: ``value``. This may be a nested structure for nested comprehensions. See
3952 #: :class:`CompFor` for details.
3953 for_in: CompFor
3954
3955 lbrace: LeftCurlyBrace = LeftCurlyBrace.field()
3956 #: Braces surrounding the dict comprehension.
3957 rbrace: RightCurlyBrace = RightCurlyBrace.field()
3958
3959 lpar: Sequence[LeftParen] = ()
3960 #: Sequence of parenthesis for precedence dictation.
3961 rpar: Sequence[RightParen] = ()
3962
3963 #: Whitespace after the key, but before the colon in ``key : value``.
3964 whitespace_before_colon: BaseParenthesizableWhitespace = SimpleWhitespace.field("")
3965 #: Whitespace after the colon, but before the value in ``key : value``.
3966 whitespace_after_colon: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
3967
3968 def _validate(self) -> None:
3969 super(DictComp, self)._validate()
3970
3971 for_in = self.for_in
3972 if (
3973 for_in.whitespace_before.empty
3974 and not self.value._safe_to_use_with_word_operator(ExpressionPosition.LEFT)
3975 ):
3976 keyword = "async" if for_in.asynchronous else "for"
3977 raise CSTValidationError(
3978 f"Must have at least one space before '{keyword}' keyword."
3979 )
3980
3981 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "DictComp":
3982 return DictComp(
3983 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
3984 lbrace=visit_required(self, "lbrace", self.lbrace, visitor),
3985 key=visit_required(self, "key", self.key, visitor),
3986 whitespace_before_colon=visit_required(
3987 self, "whitespace_before_colon", self.whitespace_before_colon, visitor
3988 ),
3989 whitespace_after_colon=visit_required(
3990 self, "whitespace_after_colon", self.whitespace_after_colon, visitor
3991 ),
3992 value=visit_required(self, "value", self.value, visitor),
3993 for_in=visit_required(self, "for_in", self.for_in, visitor),
3994 rbrace=visit_required(self, "rbrace", self.rbrace, visitor),
3995 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
3996 )
3997
3998 def _codegen_impl(self, state: CodegenState) -> None:
3999 with self._parenthesize(state), self._braceize(state):
4000 self.key._codegen(state)
4001 self.whitespace_before_colon._codegen(state)
4002 state.add_token(":")
4003 self.whitespace_after_colon._codegen(state)
4004 self.value._codegen(state)
4005 self.for_in._codegen(state)
4006
4007
4008@add_slots
4009@dataclass(frozen=True)
4010class NamedExpr(BaseExpression):
4011 """
4012 An expression that is also an assignment, such as ``x := y + z``. Affectionately
4013 known as the walrus operator, this expression allows you to make an assignment
4014 inside an expression. This greatly simplifies loops::
4015
4016 while line := read_some_line_or_none():
4017 do_thing_with_line(line)
4018 """
4019
4020 #: The target that is being assigned to.
4021 target: BaseExpression
4022
4023 #: The expression being assigned to the target.
4024 value: BaseExpression
4025
4026 #: Sequence of parenthesis for precedence dictation.
4027 lpar: Sequence[LeftParen] = ()
4028 #: Sequence of parenthesis for precedence dictation.
4029 rpar: Sequence[RightParen] = ()
4030
4031 #: Whitespace after the target, but before the walrus operator.
4032 whitespace_before_walrus: BaseParenthesizableWhitespace = SimpleWhitespace.field(
4033 " "
4034 )
4035 #: Whitespace after the walrus operator, but before the value.
4036 whitespace_after_walrus: BaseParenthesizableWhitespace = SimpleWhitespace.field(" ")
4037
4038 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "NamedExpr":
4039 return NamedExpr(
4040 lpar=visit_sequence(self, "lpar", self.lpar, visitor),
4041 target=visit_required(self, "target", self.target, visitor),
4042 whitespace_before_walrus=visit_required(
4043 self, "whitespace_before_walrus", self.whitespace_before_walrus, visitor
4044 ),
4045 whitespace_after_walrus=visit_required(
4046 self, "whitespace_after_walrus", self.whitespace_after_walrus, visitor
4047 ),
4048 value=visit_required(self, "value", self.value, visitor),
4049 rpar=visit_sequence(self, "rpar", self.rpar, visitor),
4050 )
4051
4052 def _safe_to_use_with_word_operator(self, position: ExpressionPosition) -> bool:
4053 if position == ExpressionPosition.LEFT:
4054 return len(self.rpar) > 0 or self.value._safe_to_use_with_word_operator(
4055 position
4056 )
4057 return len(self.lpar) > 0 or self.target._safe_to_use_with_word_operator(
4058 position
4059 )
4060
4061 def _codegen_impl(self, state: CodegenState) -> None:
4062 with self._parenthesize(state):
4063 self.target._codegen(state)
4064 self.whitespace_before_walrus._codegen(state)
4065 state.add_token(":=")
4066 self.whitespace_after_walrus._codegen(state)
4067 self.value._codegen(state)