Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/libcst/_nodes/base.py: 60%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
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.
6from abc import ABC, abstractmethod
7from copy import deepcopy
8from dataclasses import dataclass, field, fields, replace
9from typing import Any, cast, ClassVar, Dict, List, Mapping, Sequence, TypeVar, Union
11from libcst import CSTLogicError
12from libcst._flatten_sentinel import FlattenSentinel
13from libcst._nodes.internal import CodegenState
14from libcst._removal_sentinel import RemovalSentinel
15from libcst._type_enforce import is_value_of_type
16from libcst._types import CSTNodeT
17from libcst._visitors import CSTTransformer, CSTVisitor, CSTVisitorT
19_CSTNodeSelfT = TypeVar("_CSTNodeSelfT", bound="CSTNode")
20_EMPTY_SEQUENCE: Sequence["CSTNode"] = ()
23class CSTValidationError(SyntaxError):
24 pass
27class CSTCodegenError(SyntaxError):
28 pass
31class _ChildrenCollectionVisitor(CSTVisitor):
32 def __init__(self) -> None:
33 self.children: List[CSTNode] = []
35 def on_visit(self, node: "CSTNode") -> bool:
36 self.children.append(node)
37 return False # Don't include transitive children
40class _ChildReplacementTransformer(CSTTransformer):
41 def __init__(
42 self, old_node: "CSTNode", new_node: Union["CSTNode", RemovalSentinel]
43 ) -> None:
44 self.old_node = old_node
45 self.new_node = new_node
47 def on_visit(self, node: "CSTNode") -> bool:
48 # If the node is one we are about to replace, we shouldn't
49 # recurse down it, that would be a waste of time.
50 return node is not self.old_node
52 def on_leave(
53 self, original_node: "CSTNode", updated_node: "CSTNode"
54 ) -> Union["CSTNode", RemovalSentinel]:
55 if original_node is self.old_node:
56 return self.new_node
57 return updated_node
60class _ChildWithChangesTransformer(CSTTransformer):
61 def __init__(self, old_node: "CSTNode", changes: Mapping[str, Any]) -> None:
62 self.old_node = old_node
63 self.changes = changes
65 def on_visit(self, node: "CSTNode") -> bool:
66 # If the node is one we are about to replace, we shouldn't
67 # recurse down it, that would be a waste of time.
68 return node is not self.old_node
70 def on_leave(self, original_node: "CSTNode", updated_node: "CSTNode") -> "CSTNode":
71 if original_node is self.old_node:
72 return updated_node.with_changes(**self.changes)
73 return updated_node
76class _NOOPVisitor(CSTTransformer):
77 pass
80def _pretty_repr(value: object) -> str:
81 if not isinstance(value, str) and isinstance(value, Sequence):
82 return _pretty_repr_sequence(value)
83 else:
84 return repr(value)
87def _pretty_repr_sequence(seq: Sequence[object]) -> str:
88 if len(seq) == 0:
89 return "[]"
90 else:
91 return "\n".join(["[", *[f"{_indent(repr(el))}," for el in seq], "]"])
94def _indent(value: str) -> str:
95 return "\n".join(f" {line}" for line in value.split("\n"))
98def _clone(val: object) -> object:
99 # We can't use isinstance(val, CSTNode) here due to poor performance
100 # of isinstance checks against ABC direct subclasses. What we're trying
101 # to do here is recursively call this functionality on subclasses, but
102 # if the attribute isn't a CSTNode, fall back to copy.deepcopy.
103 try:
104 # pyre-ignore We know this might not exist, that's the point of the
105 # attribute error and try block.
106 return val.deep_clone()
107 except AttributeError:
108 return deepcopy(val)
111@dataclass(frozen=True)
112class CSTNode(ABC):
113 __slots__: ClassVar[Sequence[str]] = ()
115 def __post_init__(self) -> None:
116 # PERF: It might make more sense to move validation work into the visitor, which
117 # would allow us to avoid validating the tree when parsing a file.
118 self._validate()
120 @classmethod
121 def __init_subclass__(cls, **kwargs: Any) -> None:
122 """
123 HACK: Add our implementation of `__repr__`, `__hash__`, and `__eq__` to the
124 class's __dict__ to prevent dataclass from generating it's own `__repr__`,
125 `__hash__`, and `__eq__`.
127 The alternative is to require each implementation of a node to remember to add
128 `repr=False, eq=False`, which is more error-prone.
129 """
130 super().__init_subclass__(**kwargs)
132 if "__repr__" not in cls.__dict__:
133 cls.__repr__ = CSTNode.__repr__
134 if "__eq__" not in cls.__dict__:
135 cls.__eq__ = CSTNode.__eq__
136 if "__hash__" not in cls.__dict__:
137 cls.__hash__ = CSTNode.__hash__
139 def _validate(self) -> None:
140 """
141 Override this to perform runtime validation of a newly created node.
143 The function is called during `__init__`. It should check for possible mistakes
144 that wouldn't be caught by a static type checker.
146 If you can't use a static type checker, and want to perform a runtime validation
147 of this node's types, use `validate_types` instead.
148 """
149 pass
151 def validate_types_shallow(self) -> None:
152 """
153 Compares the type annotations on a node's fields with those field's actual
154 values at runtime. Raises a TypeError is a mismatch is found.
156 Only validates the current node, not any of it's children. For a recursive
157 version, see :func:`validate_types_deep`.
159 If you're using a static type checker (highly recommended), this is useless.
160 However, if your code doesn't use a static type checker, or if you're unable to
161 statically type your code for some reason, you can use this method to help
162 validate your tree.
164 Some (non-typing) validation is done unconditionally during the construction of
165 a node. That validation does not overlap with the work that
166 :func:`validate_types_deep` does.
167 """
168 for f in fields(self):
169 value = getattr(self, f.name)
170 if not is_value_of_type(value, f.type):
171 raise TypeError(
172 f"Expected an instance of {f.type!r} on "
173 + f"{type(self).__name__}'s '{f.name}' field, but instead got "
174 + f"an instance of {type(value)!r}"
175 )
177 def validate_types_deep(self) -> None:
178 """
179 Like :func:`validate_types_shallow`, but recursively validates the whole tree.
180 """
181 self.validate_types_shallow()
182 for ch in self.children:
183 ch.validate_types_deep()
185 @property
186 def children(self) -> Sequence["CSTNode"]:
187 """
188 The immediate (not transitive) child CSTNodes of the current node. Various
189 properties on the nodes, such as string values, will not be visited if they are
190 not a subclass of CSTNode.
192 Iterable properties of the node (e.g. an IndentedBlock's body) will be flattened
193 into the children's sequence.
195 The children will always be returned in the same order that they appear
196 lexically in the code.
197 """
199 # We're hooking into _visit_and_replace_children, which means that our current
200 # implementation is slow. We may need to rethink and/or cache this if it becomes
201 # a frequently accessed property.
202 #
203 # This probably won't be called frequently, because most child access will
204 # probably through visit, or directly through named property access, not through
205 # children.
207 visitor = _ChildrenCollectionVisitor()
208 self._visit_and_replace_children(visitor)
209 return visitor.children
211 def visit(
212 self: _CSTNodeSelfT, visitor: CSTVisitorT
213 ) -> Union[_CSTNodeSelfT, RemovalSentinel, FlattenSentinel[_CSTNodeSelfT]]:
214 """
215 Visits the current node, its children, and all transitive children using
216 the given visitor's callbacks.
217 """
218 # visit self
219 should_visit_children = visitor.on_visit(self)
221 # TODO: provide traversal where children are not replaced
222 # visit children (optionally)
223 if should_visit_children:
224 # It's not possible to define `_visit_and_replace_children` with the correct
225 # return type in any sane way, so we're using this cast. See the
226 # explanation above the declaration of `_visit_and_replace_children`.
227 with_updated_children = cast(
228 _CSTNodeSelfT, self._visit_and_replace_children(visitor)
229 )
230 else:
231 with_updated_children = self
233 if isinstance(visitor, CSTVisitor):
234 visitor.on_leave(self)
235 leave_result = self
236 else:
237 leave_result = visitor.on_leave(self, with_updated_children)
239 # validate return type of the user-defined `visitor.on_leave` method
240 if not isinstance(leave_result, (CSTNode, RemovalSentinel, FlattenSentinel)):
241 raise CSTValidationError(
242 "Expected a node of type CSTNode or a RemovalSentinel, "
243 + f"but got a return value of {type(leave_result).__name__}"
244 )
246 # TODO: Run runtime typechecks against updated nodes
248 return leave_result
250 # The return type of `_visit_and_replace_children` is `CSTNode`, not
251 # `_CSTNodeSelfT`. This is because pyre currently doesn't have a way to annotate
252 # classes as final. https://mypy.readthedocs.io/en/latest/final_attrs.html
253 #
254 # The issue is that any reasonable implementation of `_visit_and_replace_children`
255 # needs to refer to the class' own constructor:
256 #
257 # class While(CSTNode):
258 # def _visit_and_replace_children(self, visitor: CSTVisitorT) -> While:
259 # return While(...)
260 #
261 # You'll notice that because this implementation needs to call the `While`
262 # constructor, the return type is also `While`. This function is a valid subtype of
263 # `Callable[[CSTVisitorT], CSTNode]`.
264 #
265 # It is not a valid subtype of `Callable[[CSTVisitorT], _CSTNodeSelfT]`. That's
266 # because the return type of this function wouldn't be valid for any subclasses.
267 # In practice, that's not an issue, because we don't have any subclasses of `While`,
268 # but there's no way to tell pyre that without a `@final` annotation.
269 #
270 # Instead, we're just relying on an unchecked call to `cast()` in the `visit`
271 # method.
272 @abstractmethod
273 def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "CSTNode":
274 """
275 Intended to be overridden by subclasses to provide a low-level hook for the
276 visitor API.
278 Don't call this directly. Instead, use `visitor.visit_and_replace_node` or
279 `visitor.visit_and_replace_module`. If you need list of children, access the
280 `children` property instead.
282 The general expectation is that children should be visited in the order in which
283 they appear lexically.
284 """
285 ...
287 def _is_removable(self) -> bool:
288 """
289 Intended to be overridden by nodes that will be iterated over inside
290 Module and IndentedBlock. Returning true signifies that this node is
291 essentially useless and can be dropped when doing a visit across it.
292 """
293 return False
295 @abstractmethod
296 def _codegen_impl(self, state: CodegenState) -> None: ...
298 def _codegen(self, state: CodegenState, **kwargs: Any) -> None:
299 state.before_codegen(self)
300 self._codegen_impl(state, **kwargs)
301 state.after_codegen(self)
303 def with_changes(self: _CSTNodeSelfT, **changes: Any) -> _CSTNodeSelfT:
304 """
305 A convenience method for performing mutation-like operations on immutable nodes.
306 Creates a new object of the same type, replacing fields with values from the
307 supplied keyword arguments.
309 For example, to update the test of an if conditional, you could do::
311 def leave_If(self, original_node: cst.If, updated_node: cst.If) -> cst.If:
312 new_node = updated_node.with_changes(test=new_conditional)
313 return new_node
315 ``new_node`` will have the same ``body``, ``orelse``, and whitespace fields as
316 ``updated_node``, but with the updated ``test`` field.
318 The accepted arguments match the arguments given to ``__init__``, however there
319 are no required or positional arguments.
321 TODO: This API is untyped. There's probably no sane way to type it using pyre's
322 current feature-set, but we should still think about ways to type this or a
323 similar API in the future.
324 """
325 return replace(self, **changes)
327 def deep_clone(self: _CSTNodeSelfT) -> _CSTNodeSelfT:
328 """
329 Recursively clone the entire tree. The created tree is a new tree has the same
330 representation but different identity.
332 >>> tree = cst.parse_expression("1+2")
334 >>> tree.deep_clone() == tree
335 False
337 >>> tree == tree
338 True
340 >>> tree.deep_equals(tree.deep_clone())
341 True
342 """
343 cloned_fields: Dict[str, object] = {}
344 for field in fields(self):
345 key = field.name
346 if key[0] == "_":
347 continue
348 val = getattr(self, key)
350 # Much like the comment on _clone itself, we are allergic to instance
351 # checks against Sequence because of speed issues with ABC classes. So,
352 # instead, first handle sequence types that we do not want to iterate on
353 # and then just try to iterate and clone.
354 if isinstance(val, (str, bytes)):
355 cloned_fields[key] = _clone(val)
356 else:
357 try:
358 cloned_fields[key] = tuple(_clone(v) for v in val)
359 except TypeError:
360 cloned_fields[key] = _clone(val)
362 return type(self)(**cloned_fields)
364 def deep_equals(self, other: "CSTNode") -> bool:
365 """
366 Recursively inspects the entire tree under ``self`` and ``other`` to determine if
367 the two trees are equal by representation instead of identity (``==``).
368 """
369 from libcst._nodes.deep_equals import deep_equals as deep_equals_impl
371 return deep_equals_impl(self, other)
373 def deep_replace(
374 self: _CSTNodeSelfT, old_node: "CSTNode", new_node: CSTNodeT
375 ) -> Union[_CSTNodeSelfT, CSTNodeT]:
376 """
377 Recursively replaces any instance of ``old_node`` with ``new_node`` by identity.
378 Use this to avoid nested ``with_changes`` blocks when you are replacing one of
379 a node's deep children with a new node. Note that if you have previously
380 modified the tree in a way that ``old_node`` appears more than once as a deep
381 child, all instances will be replaced.
382 """
383 new_tree = self.visit(_ChildReplacementTransformer(old_node, new_node))
384 if isinstance(new_tree, (FlattenSentinel, RemovalSentinel)):
385 # The above transform never returns *Sentinel, so this isn't possible
386 raise CSTLogicError("Logic error, cannot get a *Sentinel here!")
387 return new_tree
389 def deep_remove(
390 self: _CSTNodeSelfT, old_node: "CSTNode"
391 ) -> Union[_CSTNodeSelfT, RemovalSentinel]:
392 """
393 Recursively removes any instance of ``old_node`` by identity. Note that if you
394 have previously modified the tree in a way that ``old_node`` appears more than
395 once as a deep child, all instances will be removed.
396 """
397 new_tree = self.visit(
398 _ChildReplacementTransformer(old_node, RemovalSentinel.REMOVE)
399 )
401 if isinstance(new_tree, FlattenSentinel):
402 # The above transform never returns FlattenSentinel, so this isn't possible
403 raise CSTLogicError("Logic error, cannot get a FlattenSentinel here!")
405 return new_tree
407 def with_deep_changes(
408 self: _CSTNodeSelfT, old_node: "CSTNode", **changes: Any
409 ) -> _CSTNodeSelfT:
410 """
411 A convenience method for applying :attr:`with_changes` to a child node. Use
412 this to avoid chains of :attr:`with_changes` or combinations of
413 :attr:`deep_replace` and :attr:`with_changes`.
415 The accepted arguments match the arguments given to the child node's
416 ``__init__``.
418 TODO: This API is untyped. There's probably no sane way to type it using pyre's
419 current feature-set, but we should still think about ways to type this or a
420 similar API in the future.
421 """
422 new_tree = self.visit(_ChildWithChangesTransformer(old_node, changes))
423 if isinstance(new_tree, (FlattenSentinel, RemovalSentinel)):
424 # This is impossible with the above transform.
425 raise CSTLogicError("Logic error, cannot get a *Sentinel here!")
426 return new_tree
428 def __eq__(self: _CSTNodeSelfT, other: object) -> bool:
429 """
430 CSTNodes are only treated as equal by identity. This matches the behavior of
431 CPython's AST nodes.
433 If you actually want to compare the value instead of the identity of the current
434 node with another, use `node.deep_equals`. Because `deep_equals` must traverse
435 the entire tree, it can have an unexpectedly large time complexity.
437 We're not exposing value equality as the default behavior because of
438 `deep_equals`'s large time complexity.
439 """
440 return self is other
442 def __hash__(self) -> int:
443 # Equality of nodes is based on identity, so the hash should be too.
444 return id(self)
446 def __repr__(self) -> str:
447 if len(fields(self)) == 0:
448 return f"{type(self).__name__}()"
450 lines = [f"{type(self).__name__}("]
451 for f in fields(self):
452 key = f.name
453 if key[0] != "_":
454 value = getattr(self, key)
455 lines.append(_indent(f"{key}={_pretty_repr(value)},"))
456 lines.append(")")
457 return "\n".join(lines)
459 @classmethod
460 # pyre-fixme[3]: Return annotation cannot be `Any`.
461 def field(cls, *args: object, **kwargs: object) -> Any:
462 """
463 A helper that allows us to easily use CSTNodes in dataclass constructor
464 defaults without accidentally aliasing nodes by identity across multiple
465 instances.
466 """
467 # pyre-ignore Pyre is complaining about CSTNode not being instantiable,
468 # but we're only going to call this from concrete subclasses.
469 return field(default_factory=lambda: cls(*args, **kwargs))
472class BaseLeaf(CSTNode, ABC):
473 __slots__ = ()
475 @property
476 def children(self) -> Sequence[CSTNode]:
477 # override this with an optimized implementation
478 return _EMPTY_SEQUENCE
480 def _visit_and_replace_children(
481 self: _CSTNodeSelfT, visitor: CSTVisitorT
482 ) -> _CSTNodeSelfT:
483 return self
486class BaseValueToken(BaseLeaf, ABC):
487 """
488 Represents the subset of nodes that only contain a value. Not all tokens from the
489 tokenizer will exist as BaseValueTokens. In places where the token is always a
490 constant value (e.g. a COLON token), the token's value will be implicitly folded
491 into the parent CSTNode, and hard-coded into the implementation of _codegen.
492 """
494 __slots__ = ()
496 value: str
498 def _codegen_impl(self, state: CodegenState) -> None:
499 state.add_token(self.value)