visitor.py

from copy import copy
from enum import Enum
from typing import (
    Any,
    Callable,
    Collection,
    Dict,
    List,
    NamedTuple,
    Optional,
    Tuple,
    Union,
)

from ..pyutils import inspect, snake_to_camel
from . import ast

from .ast import Node, QUERY_DOCUMENT_KEYS

__all__ = [
    "Visitor",
    "ParallelVisitor",
    "VisitorAction",
    "visit",
    "BREAK",
    "SKIP",
    "REMOVE",
    "IDLE",
]


class VisitorActionEnum(Enum):
    """Special return values for the visitor methods.

    You can also use the values of this enum directly.
    """

    BREAK = True
    SKIP = False
    REMOVE = Ellipsis


VisitorAction = Optional[VisitorActionEnum]

# Note that in GraphQL.js these are defined differently:
# BREAK = {}, SKIP = false, REMOVE = null, IDLE = undefined

BREAK = VisitorActionEnum.BREAK
SKIP = VisitorActionEnum.SKIP
REMOVE = VisitorActionEnum.REMOVE
IDLE = None

VisitorKeyMap = Dict[str, Tuple[str, ...]]


class EnterLeaveVisitor(NamedTuple):
    """Visitor with functions for entering and leaving."""

    enter: Optional[Callable[..., Optional[VisitorAction]]]
    leave: Optional[Callable[..., Optional[VisitorAction]]]


class Visitor:
    """Visitor that walks through an AST.

    Visitors can define two generic methods "enter" and "leave". The former will be
    called when a node is entered in the traversal, the latter is called after visiting
    the node and its child nodes. These methods have the following signature::

        def enter(self, node, key, parent, path, ancestors):
            # The return value has the following meaning:
            # IDLE (None): no action
            # SKIP: skip visiting this node
            # BREAK: stop visiting altogether
            # REMOVE: delete this node
            # any other value: replace this node with the returned value
            return

        def leave(self, node, key, parent, path, ancestors):
            # The return value has the following meaning:
            # IDLE (None) or SKIP: no action
            # BREAK: stop visiting altogether
            # REMOVE: delete this node
            # any other value: replace this node with the returned value
            return

    The parameters have the following meaning:

    :arg node: The current node being visiting.
    :arg key: The index or key to this node from the parent node or Array.
    :arg parent: the parent immediately above this node, which may be an Array.
    :arg path: The key path to get to this node from the root node.
    :arg ancestors: All nodes and Arrays visited before reaching parent
        of this node. These correspond to array indices in ``path``.
        Note: ancestors includes arrays which contain the parent of visited node.

    You can also define node kind specific methods by suffixing them with an underscore
    followed by the kind of the node to be visited. For instance, to visit ``field``
    nodes, you would defined the methods ``enter_field()`` and/or ``leave_field()``,
    with the same signature as above. If no kind specific method has been defined
    for a given node, the generic method is called.
    """

    # Provide special return values as attributes
    BREAK, SKIP, REMOVE, IDLE = BREAK, SKIP, REMOVE, IDLE

    enter_leave_map: Dict[str, EnterLeaveVisitor]

    def __init_subclass__(cls) -> None:
        """Verify that all defined handlers are valid."""
        super().__init_subclass__()
        for attr, val in cls.__dict__.items():
            if attr.startswith("_"):
                continue
            attr_kind = attr.split("_", 1)
            if len(attr_kind) < 2:
                kind: Optional[str] = None
            else:
                attr, kind = attr_kind
            if attr in ("enter", "leave") and kind:
                name = snake_to_camel(kind) + "Node"
                node_cls = getattr(ast, name, None)
                if (
                    not node_cls
                    or not isinstance(node_cls, type)
                    or not issubclass(node_cls, Node)
                ):
                    raise TypeError(f"Invalid AST node kind: {kind}.")

    def __init__(self) -> None:
        self.enter_leave_map = {}

    def get_enter_leave_for_kind(self, kind: str) -> EnterLeaveVisitor:
        """Given a node kind, return the EnterLeaveVisitor for that kind."""
        try:
            return self.enter_leave_map[kind]
        except KeyError:
            enter_fn = getattr(self, f"enter_{kind}", None)
            if not enter_fn:
                enter_fn = getattr(self, "enter", None)
            leave_fn = getattr(self, f"leave_{kind}", None)
            if not leave_fn:
                leave_fn = getattr(self, "leave", None)
            enter_leave = EnterLeaveVisitor(enter_fn, leave_fn)
            self.enter_leave_map[kind] = enter_leave
            return enter_leave

    def get_visit_fn(
        self, kind: str, is_leaving: bool = False
    ) -> Optional[Callable[..., Optional[VisitorAction]]]:
        """Get the visit function for the given node kind and direction.

        .. deprecated:: 3.2
           Please use ``get_enter_leave_for_kind`` instead. Will be removed in v3.3.
        """
        enter_leave = self.get_enter_leave_for_kind(kind)
        return enter_leave.leave if is_leaving else enter_leave.enter


class Stack(NamedTuple):
    """A stack for the visit function."""

    in_array: bool
    idx: int
    keys: Tuple[Node, ...]
    edits: List[Tuple[Union[int, str], Node]]
    prev: Any  # 'Stack' (python/mypy/issues/731)


def visit(
    root: Node, visitor: Visitor, visitor_keys: Optional[VisitorKeyMap] = None
) -> Any:
    """Visit each node in an AST.

    :func:`~.visit` will walk through an AST using a depth-first traversal, calling the
    visitor's enter methods at each node in the traversal, and calling the leave methods
    after visiting that node and all of its child nodes.

    By returning different values from the enter and leave methods, the behavior of the
    visitor can be altered, including skipping over a sub-tree of the AST (by returning
    False), editing the AST by returning a value or None to remove the value, or to stop
    the whole traversal by returning :data:`~.BREAK`.

    When using :func:`~.visit` to edit an AST, the original AST will not be modified,
    and a new version of the AST with the changes applied will be returned from the
    visit function.

    To customize the node attributes to be used for traversal, you can provide a
    dictionary visitor_keys mapping node kinds to node attributes.
    """
    if not isinstance(root, Node):
        raise TypeError(f"Not an AST Node: {inspect(root)}.")
    if not isinstance(visitor, Visitor):
        raise TypeError(f"Not an AST Visitor: {inspect(visitor)}.")
    if visitor_keys is None:
        visitor_keys = QUERY_DOCUMENT_KEYS

    stack: Any = None
    in_array = False
    keys: Tuple[Node, ...] = (root,)
    idx = -1
    edits: List[Any] = []
    node: Any = root
    key: Any = None
    parent: Any = None
    path: List[Any] = []
    path_append = path.append
    path_pop = path.pop
    ancestors: List[Any] = []
    ancestors_append = ancestors.append
    ancestors_pop = ancestors.pop

    while True:
        idx += 1
        is_leaving = idx == len(keys)
        is_edited = is_leaving and edits
        if is_leaving:
            key = path[-1] if ancestors else None
            node = parent
            parent = ancestors_pop() if ancestors else None
            if is_edited:
                if in_array:
                    node = list(node)
                    edit_offset = 0
                    for edit_key, edit_value in edits:
                        array_key = edit_key - edit_offset
                        if edit_value is REMOVE or edit_value is Ellipsis:
                            node.pop(array_key)
                            edit_offset += 1
                        else:
                            node[array_key] = edit_value
                    node = tuple(node)
                else:
                    node = copy(node)
                    for edit_key, edit_value in edits:
                        setattr(node, edit_key, edit_value)
            idx = stack.idx
            keys = stack.keys
            edits = stack.edits
            in_array = stack.in_array
            stack = stack.prev
        elif parent:
            if in_array:
                key = idx
                node = parent[key]
            else:
                key = keys[idx]
                node = getattr(parent, key, None)
            if node is None:
                continue
            path_append(key)

        if isinstance(node, tuple):
            result = None
        else:
            if not isinstance(node, Node):
                raise TypeError(f"Invalid AST Node: {inspect(node)}.")
            enter_leave = visitor.get_enter_leave_for_kind(node.kind)
            visit_fn = enter_leave.leave if is_leaving else enter_leave.enter
            if visit_fn:
                result = visit_fn(node, key, parent, path, ancestors)

                if result is BREAK or result is True:
                    break

                if result is SKIP or result is False:
                    if not is_leaving:
                        path_pop()
                        continue

                elif result is not None:
                    edits.append((key, result))
                    if not is_leaving:
                        if isinstance(result, Node):
                            node = result
                        else:
                            path_pop()
                            continue
            else:
                result = None

        if result is None and is_edited:
            edits.append((key, node))

        if is_leaving:
            if path:
                path_pop()
        else:
            stack = Stack(in_array, idx, keys, edits, stack)
            in_array = isinstance(node, tuple)
            keys = node if in_array else visitor_keys.get(node.kind, ())
            idx = -1
            edits = []
            if parent:
                ancestors_append(parent)
            parent = node

        if not stack:
            break

    if edits:
        return edits[-1][1]

    return root


class ParallelVisitor(Visitor):
    """A Visitor which delegates to many visitors to run in parallel.

    Each visitor will be visited for each node before moving on.

    If a prior visitor edits a node, no following visitors will see that node.
    """

    def __init__(self, visitors: Collection[Visitor]):
        """Create a new visitor from the given list of parallel visitors."""
        super().__init__()
        self.visitors = visitors
        self.skipping: List[Any] = [None] * len(visitors)

    def get_enter_leave_for_kind(self, kind: str) -> EnterLeaveVisitor:
        """Given a node kind, return the EnterLeaveVisitor for that kind."""
        try:
            return self.enter_leave_map[kind]
        except KeyError:
            has_visitor = False
            enter_list: List[Optional[Callable[..., Optional[VisitorAction]]]] = []
            leave_list: List[Optional[Callable[..., Optional[VisitorAction]]]] = []
            for visitor in self.visitors:
                enter, leave = visitor.get_enter_leave_for_kind(kind)
                if not has_visitor and (enter or leave):
                    has_visitor = True
                enter_list.append(enter)
                leave_list.append(leave)

            if has_visitor:

                def enter(node: Node, *args: Any) -> Optional[VisitorAction]:
                    skipping = self.skipping
                    for i, fn in enumerate(enter_list):
                        if not skipping[i]:
                            if fn:
                                result = fn(node, *args)
                                if result is SKIP or result is False:
                                    skipping[i] = node
                                elif result is BREAK or result is True:
                                    skipping[i] = BREAK
                                elif result is not None:
                                    return result
                    return None

                def leave(node: Node, *args: Any) -> Optional[VisitorAction]:
                    skipping = self.skipping
                    for i, fn in enumerate(leave_list):
                        if not skipping[i]:
                            if fn:
                                result = fn(node, *args)
                                if result is BREAK or result is True:
                                    skipping[i] = BREAK
                                elif (
                                    result is not None
                                    and result is not SKIP
                                    and result is not False
                                ):
                                    return result
                        elif skipping[i] is node:
                            skipping[i] = None
                    return None

            else:

                enter = leave = None

            enter_leave = EnterLeaveVisitor(enter, leave)
            self.enter_leave_map[kind] = enter_leave
            return enter_leave