base.py

# 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)


@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

    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}"
                )

    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

    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)

    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)

    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):  # noqa: F402
            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)

    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)

    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

    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

    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)

    @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)