# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from abc import ABC, abstractmethod
from copy import deepcopy
from dataclasses import dataclass, field, fields, replace
from typing import Any, cast, ClassVar, Dict, List, Mapping, Sequence, TypeVar, Union
from libcst._flatten_sentinel import FlattenSentinel
from libcst._nodes.internal import CodegenState
from libcst._removal_sentinel import RemovalSentinel
from libcst._type_enforce import is_value_of_type
from libcst._types import CSTNodeT
from libcst._visitors import CSTTransformer, CSTVisitor, CSTVisitorT
_CSTNodeSelfT = TypeVar("_CSTNodeSelfT", bound="CSTNode")
_EMPTY_SEQUENCE: Sequence["CSTNode"] = ()
class CSTValidationError(SyntaxError):
pass
class CSTCodegenError(SyntaxError):
pass
class _ChildrenCollectionVisitor(CSTVisitor):
def __init__(self) -> None:
self.children: List[CSTNode] = []
def on_visit(self, node: "CSTNode") -> bool:
self.children.append(node)
return False # Don't include transitive children
class _ChildReplacementTransformer(CSTTransformer):
def __init__(
self, old_node: "CSTNode", new_node: Union["CSTNode", RemovalSentinel]
) -> None:
self.old_node = old_node
self.new_node = new_node
def on_visit(self, node: "CSTNode") -> bool:
# If the node is one we are about to replace, we shouldn't
# recurse down it, that would be a waste of time.
return node is not self.old_node
def on_leave(
self, original_node: "CSTNode", updated_node: "CSTNode"
) -> Union["CSTNode", RemovalSentinel]:
if original_node is self.old_node:
return self.new_node
return updated_node
class _ChildWithChangesTransformer(CSTTransformer):
def __init__(self, old_node: "CSTNode", changes: Mapping[str, Any]) -> None:
self.old_node = old_node
self.changes = changes
def on_visit(self, node: "CSTNode") -> bool:
# If the node is one we are about to replace, we shouldn't
# recurse down it, that would be a waste of time.
return node is not self.old_node
def on_leave(self, original_node: "CSTNode", updated_node: "CSTNode") -> "CSTNode":
if original_node is self.old_node:
return updated_node.with_changes(**self.changes)
return updated_node
class _NOOPVisitor(CSTTransformer):
pass
def _pretty_repr(value: object) -> str:
if not isinstance(value, str) and isinstance(value, Sequence):
return _pretty_repr_sequence(value)
else:
return repr(value)
def _pretty_repr_sequence(seq: Sequence[object]) -> str:
if len(seq) == 0:
return "[]"
else:
return "\n".join(["[", *[f"{_indent(repr(el))}," for el in seq], "]"])
def _indent(value: str) -> str:
return "\n".join(f" {line}" for line in value.split("\n"))
def _clone(val: object) -> object:
# We can't use isinstance(val, CSTNode) here due to poor performance
# of isinstance checks against ABC direct subclasses. What we're trying
# to do here is recursively call this functionality on subclasses, but
# if the attribute isn't a CSTNode, fall back to copy.deepcopy.
try:
# pyre-ignore We know this might not exist, that's the point of the
# attribute error and try block.
return val.deep_clone()
except AttributeError:
return deepcopy(val)
[docs]
@dataclass(frozen=True)
class CSTNode(ABC):
__slots__: ClassVar[Sequence[str]] = ()
def __post_init__(self) -> None:
# PERF: It might make more sense to move validation work into the visitor, which
# would allow us to avoid validating the tree when parsing a file.
self._validate()
@classmethod
def __init_subclass__(cls, **kwargs: Any) -> None:
"""
HACK: Add our implementation of `__repr__`, `__hash__`, and `__eq__` to the
class's __dict__ to prevent dataclass from generating it's own `__repr__`,
`__hash__`, and `__eq__`.
The alternative is to require each implementation of a node to remember to add
`repr=False, eq=False`, which is more error-prone.
"""
super().__init_subclass__(**kwargs)
if "__repr__" not in cls.__dict__:
cls.__repr__ = CSTNode.__repr__
if "__eq__" not in cls.__dict__:
cls.__eq__ = CSTNode.__eq__
if "__hash__" not in cls.__dict__:
cls.__hash__ = CSTNode.__hash__
def _validate(self) -> None:
"""
Override this to perform runtime validation of a newly created node.
The function is called during `__init__`. It should check for possible mistakes
that wouldn't be caught by a static type checker.
If you can't use a static type checker, and want to perform a runtime validation
of this node's types, use `validate_types` instead.
"""
pass
[docs]
def validate_types_shallow(self) -> None:
"""
Compares the type annotations on a node's fields with those field's actual
values at runtime. Raises a TypeError is a mismatch is found.
Only validates the current node, not any of it's children. For a recursive
version, see :func:`validate_types_deep`.
If you're using a static type checker (highly recommended), this is useless.
However, if your code doesn't use a static type checker, or if you're unable to
statically type your code for some reason, you can use this method to help
validate your tree.
Some (non-typing) validation is done unconditionally during the construction of
a node. That validation does not overlap with the work that
:func:`validate_types_deep` does.
"""
for f in fields(self):
value = getattr(self, f.name)
if not is_value_of_type(value, f.type):
raise TypeError(
f"Expected an instance of {f.type!r} on "
+ f"{type(self).__name__}'s '{f.name}' field, but instead got "
+ f"an instance of {type(value)!r}"
)
[docs]
def validate_types_deep(self) -> None:
"""
Like :func:`validate_types_shallow`, but recursively validates the whole tree.
"""
self.validate_types_shallow()
for ch in self.children:
ch.validate_types_deep()
@property
def children(self) -> Sequence["CSTNode"]:
"""
The immediate (not transitive) child CSTNodes of the current node. Various
properties on the nodes, such as string values, will not be visited if they are
not a subclass of CSTNode.
Iterable properties of the node (e.g. an IndentedBlock's body) will be flattened
into the children's sequence.
The children will always be returned in the same order that they appear
lexically in the code.
"""
# We're hooking into _visit_and_replace_children, which means that our current
# implementation is slow. We may need to rethink and/or cache this if it becomes
# a frequently accessed property.
#
# This probably won't be called frequently, because most child access will
# probably through visit, or directly through named property access, not through
# children.
visitor = _ChildrenCollectionVisitor()
self._visit_and_replace_children(visitor)
return visitor.children
[docs]
def visit(
self: _CSTNodeSelfT, visitor: CSTVisitorT
) -> Union[_CSTNodeSelfT, RemovalSentinel, FlattenSentinel[_CSTNodeSelfT]]:
"""
Visits the current node, its children, and all transitive children using
the given visitor's callbacks.
"""
# visit self
should_visit_children = visitor.on_visit(self)
# TODO: provide traversal where children are not replaced
# visit children (optionally)
if should_visit_children:
# It's not possible to define `_visit_and_replace_children` with the correct
# return type in any sane way, so we're using this cast. See the
# explanation above the declaration of `_visit_and_replace_children`.
with_updated_children = cast(
_CSTNodeSelfT, self._visit_and_replace_children(visitor)
)
else:
with_updated_children = self
if isinstance(visitor, CSTVisitor):
visitor.on_leave(self)
leave_result = self
else:
leave_result = visitor.on_leave(self, with_updated_children)
# validate return type of the user-defined `visitor.on_leave` method
if not isinstance(leave_result, (CSTNode, RemovalSentinel, FlattenSentinel)):
raise Exception(
"Expected a node of type CSTNode or a RemovalSentinel, "
+ f"but got a return value of {type(leave_result).__name__}"
)
# TODO: Run runtime typechecks against updated nodes
return leave_result
# The return type of `_visit_and_replace_children` is `CSTNode`, not
# `_CSTNodeSelfT`. This is because pyre currently doesn't have a way to annotate
# classes as final. https://mypy.readthedocs.io/en/latest/final_attrs.html
#
# The issue is that any reasonable implementation of `_visit_and_replace_children`
# needs to refer to the class' own constructor:
#
# class While(CSTNode):
# def _visit_and_replace_children(self, visitor: CSTVisitorT) -> While:
# return While(...)
#
# You'll notice that because this implementation needs to call the `While`
# constructor, the return type is also `While`. This function is a valid subtype of
# `Callable[[CSTVisitorT], CSTNode]`.
#
# It is not a valid subtype of `Callable[[CSTVisitorT], _CSTNodeSelfT]`. That's
# because the return type of this function wouldn't be valid for any subclasses.
# In practice, that's not an issue, because we don't have any subclasses of `While`,
# but there's no way to tell pyre that without a `@final` annotation.
#
# Instead, we're just relying on an unchecked call to `cast()` in the `visit`
# method.
@abstractmethod
def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "CSTNode":
"""
Intended to be overridden by subclasses to provide a low-level hook for the
visitor API.
Don't call this directly. Instead, use `visitor.visit_and_replace_node` or
`visitor.visit_and_replace_module`. If you need list of children, access the
`children` property instead.
The general expectation is that children should be visited in the order in which
they appear lexically.
"""
...
def _is_removable(self) -> bool:
"""
Intended to be overridden by nodes that will be iterated over inside
Module and IndentedBlock. Returning true signifies that this node is
essentially useless and can be dropped when doing a visit across it.
"""
return False
@abstractmethod
def _codegen_impl(self, state: CodegenState) -> None:
...
def _codegen(self, state: CodegenState, **kwargs: Any) -> None:
state.before_codegen(self)
self._codegen_impl(state, **kwargs)
state.after_codegen(self)
[docs]
def with_changes(self: _CSTNodeSelfT, **changes: Any) -> _CSTNodeSelfT:
"""
A convenience method for performing mutation-like operations on immutable nodes.
Creates a new object of the same type, replacing fields with values from the
supplied keyword arguments.
For example, to update the test of an if conditional, you could do::
def leave_If(self, original_node: cst.If, updated_node: cst.If) -> cst.If:
new_node = updated_node.with_changes(test=new_conditional)
return new_node
``new_node`` will have the same ``body``, ``orelse``, and whitespace fields as
``updated_node``, but with the updated ``test`` field.
The accepted arguments match the arguments given to ``__init__``, however there
are no required or positional arguments.
TODO: This API is untyped. There's probably no sane way to type it using pyre's
current feature-set, but we should still think about ways to type this or a
similar API in the future.
"""
return replace(self, **changes)
[docs]
def deep_clone(self: _CSTNodeSelfT) -> _CSTNodeSelfT:
"""
Recursively clone the entire tree. The created tree is a new tree has the same
representation but different identity.
>>> tree = cst.parse_expression("1+2")
>>> tree.deep_clone() == tree
False
>>> tree == tree
True
>>> tree.deep_equals(tree.deep_clone())
True
"""
cloned_fields: Dict[str, object] = {}
for field in fields(self):
key = field.name
if key[0] == "_":
continue
val = getattr(self, key)
# Much like the comment on _clone itself, we are allergic to instance
# checks against Sequence because of speed issues with ABC classes. So,
# instead, first handle sequence types that we do not want to iterate on
# and then just try to iterate and clone.
if isinstance(val, (str, bytes)):
cloned_fields[key] = _clone(val)
else:
try:
cloned_fields[key] = tuple(_clone(v) for v in val)
except TypeError:
cloned_fields[key] = _clone(val)
return type(self)(**cloned_fields)
[docs]
def deep_equals(self, other: "CSTNode") -> bool:
"""
Recursively inspects the entire tree under ``self`` and ``other`` to determine if
the two trees are equal by representation instead of identity (``==``).
"""
from libcst._nodes.deep_equals import deep_equals as deep_equals_impl
return deep_equals_impl(self, other)
[docs]
def deep_replace(
self: _CSTNodeSelfT, old_node: "CSTNode", new_node: CSTNodeT
) -> Union[_CSTNodeSelfT, CSTNodeT]:
"""
Recursively replaces any instance of ``old_node`` with ``new_node`` by identity.
Use this to avoid nested ``with_changes`` blocks when you are replacing one of
a node's deep children with a new node. Note that if you have previously
modified the tree in a way that ``old_node`` appears more than once as a deep
child, all instances will be replaced.
"""
new_tree = self.visit(_ChildReplacementTransformer(old_node, new_node))
if isinstance(new_tree, (FlattenSentinel, RemovalSentinel)):
# The above transform never returns *Sentinel, so this isn't possible
raise Exception("Logic error, cannot get a *Sentinel here!")
return new_tree
[docs]
def deep_remove(
self: _CSTNodeSelfT, old_node: "CSTNode"
) -> Union[_CSTNodeSelfT, RemovalSentinel]:
"""
Recursively removes any instance of ``old_node`` by identity. Note that if you
have previously modified the tree in a way that ``old_node`` appears more than
once as a deep child, all instances will be removed.
"""
new_tree = self.visit(
_ChildReplacementTransformer(old_node, RemovalSentinel.REMOVE)
)
if isinstance(new_tree, FlattenSentinel):
# The above transform never returns FlattenSentinel, so this isn't possible
raise Exception("Logic error, cannot get a FlattenSentinel here!")
return new_tree
[docs]
def with_deep_changes(
self: _CSTNodeSelfT, old_node: "CSTNode", **changes: Any
) -> _CSTNodeSelfT:
"""
A convenience method for applying :attr:`with_changes` to a child node. Use
this to avoid chains of :attr:`with_changes` or combinations of
:attr:`deep_replace` and :attr:`with_changes`.
The accepted arguments match the arguments given to the child node's
``__init__``.
TODO: This API is untyped. There's probably no sane way to type it using pyre's
current feature-set, but we should still think about ways to type this or a
similar API in the future.
"""
new_tree = self.visit(_ChildWithChangesTransformer(old_node, changes))
if isinstance(new_tree, (FlattenSentinel, RemovalSentinel)):
# This is impossible with the above transform.
raise Exception("Logic error, cannot get a *Sentinel here!")
return new_tree
def __eq__(self: _CSTNodeSelfT, other: object) -> bool:
"""
CSTNodes are only treated as equal by identity. This matches the behavior of
CPython's AST nodes.
If you actually want to compare the value instead of the identity of the current
node with another, use `node.deep_equals`. Because `deep_equals` must traverse
the entire tree, it can have an unexpectedly large time complexity.
We're not exposing value equality as the default behavior because of
`deep_equals`'s large time complexity.
"""
return self is other
def __hash__(self) -> int:
# Equality of nodes is based on identity, so the hash should be too.
return id(self)
def __repr__(self) -> str:
if len(fields(self)) == 0:
return f"{type(self).__name__}()"
lines = [f"{type(self).__name__}("]
for f in fields(self):
key = f.name
if key[0] != "_":
value = getattr(self, key)
lines.append(_indent(f"{key}={_pretty_repr(value)},"))
lines.append(")")
return "\n".join(lines)
[docs]
@classmethod
# pyre-fixme[3]: Return annotation cannot be `Any`.
def field(cls, *args: object, **kwargs: object) -> Any:
"""
A helper that allows us to easily use CSTNodes in dataclass constructor
defaults without accidentally aliasing nodes by identity across multiple
instances.
"""
# pyre-ignore Pyre is complaining about CSTNode not being instantiable,
# but we're only going to call this from concrete subclasses.
return field(default_factory=lambda: cls(*args, **kwargs))
class BaseLeaf(CSTNode, ABC):
__slots__ = ()
@property
def children(self) -> Sequence[CSTNode]:
# override this with an optimized implementation
return _EMPTY_SEQUENCE
def _visit_and_replace_children(
self: _CSTNodeSelfT, visitor: CSTVisitorT
) -> _CSTNodeSelfT:
return self
class BaseValueToken(BaseLeaf, ABC):
"""
Represents the subset of nodes that only contain a value. Not all tokens from the
tokenizer will exist as BaseValueTokens. In places where the token is always a
constant value (e.g. a COLON token), the token's value will be implicitly folded
into the parent CSTNode, and hard-coded into the implementation of _codegen.
"""
__slots__ = ()
value: str
def _codegen_impl(self, state: CodegenState) -> None:
state.add_token(self.value)