orchestrata_ultimate.refactor.ctok.node_index

Reference › orchestrata_ultimate › refactor › ctok › node_index

NodeIndex connects the ast view of a file to its byte-level tokens, so that an edit can be aimed at the exact characters of a syntax node. It is built over a TokenStream by build(tree, stream) and used by lossless_module (REQ-110-123).

Construction converts the line and column offsets that ast reports into byte offsets using a table of line starts, and indexes the tokens by start and end offset. span_for(node) returns a Token covering the node's full source range, tagged with the kind of its first token. token_for_name(node) returns the single name token that is a rename target: the identifier for a Name, the trailing attribute name for an Attribute, and for a function, async function or class the name after the leading keywords (it skips both async and def). A node without location data, or whose start does not fall on a token boundary, raises NodeSpanError; an unsupported node type raises TypeError.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

NodeSpanError

Bases: ValueError · Decorators: traces(spec='110', story='US-110-009', reqs=['REQ-110-123']) · Line 48

Raised when an ast node has no locatable token span.

NodeIndex (dataclass)

Decorators: traces(spec='110', story='US-110-009', reqs=['REQ-110-123']), dataclass(frozen=True, slots=True) · Line 107

Maps ast nodes to their token byte-spans over a TokenStream.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/refactor/ctok/node_index.py

"""NodeIndex -- ast node <-> token byte-span mapping (P2 of ctok-parser).

P2 of the ctok-parser contract (plans/contracts/ctok-parser.contract.json):
``NodeIndex`` closes the gap between P1's flat token array and CPython's
``ast`` module, so rename/edit tooling can locate the exact bytes an ast
node covers without libcst. ast node location (``lineno``/``col_offset``/
``end_lineno``/``end_col_offset``) is turned into an absolute byte offset
via a per-line offset table, then resolved against the real token array
built by :class:`~orchestrata_ultimate.refactor.ctok.tokenizer.CTokenizer`.

Two CPython quirks the offset table accounts for (both proven empirically
against ``ast.parse``, see the P2 build notes):

* ``ast`` column offsets are counted AFTER a leading UTF-8 BOM is
  stripped, while the C tokenizer's byte spans are over the RAW bytes
  (BOM included) -- the table's first line starts after the BOM so both
  line up on the same byte.
* ``ast`` treats a lone ``\\r`` exactly like ``\\r\\n``/``\\n`` as a line
  terminator (CPython applies universal-newline decoding before parsing),
  so line boundaries are found by matching any of the three, not just
  ``\\n``.
"""

from __future__ import annotations

import ast
import keyword
import re
from dataclasses import dataclass

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.refactor.ctok.token_stream import TokenStream
from orchestrata_ultimate.refactor.ctok.tokenizer import Token, TokenKind

_LINE_TERMINATOR_RE = re.compile(rb"\r\n|\r|\n")
_UTF8_BOM = b"\xef\xbb\xbf"

# ast node types whose rename target is an identifier introduced AFTER one
# or more leading keyword tokens ("class", "def", or "async" + "def").
_KEYWORD_NAMED_TYPES: tuple[type[ast.AST], ...] = (
    ast.FunctionDef,
    ast.AsyncFunctionDef,
    ast.ClassDef,
)


@traces(spec="110", story="US-110-009", reqs=["REQ-110-123"])
class NodeSpanError(ValueError):
    """Raised when an ast node has no locatable token span."""


@traces(spec="110", story="US-110-009", reqs=["REQ-110-123"])
def _location(node: ast.AST) -> tuple[int, int, int, int]:
    """Extract ``(lineno, col_offset, end_lineno, end_col_offset)``.

    Uses ``getattr`` rather than direct attribute access on purpose: only
    SOME ``ast.AST`` subclasses carry location (it lives in each node
    type's own ``_attributes``, not on the ``ast.AST`` base), and
    ``end_lineno``/``end_col_offset`` are typed ``int | None`` even on the
    nodes that do -- this is the one place that validates all four are
    real ints before any arithmetic touches them.

    Args:
        node: Any ast node.

    Returns:
        The four location ints.

    Raises:
        NodeSpanError: ``node`` is missing full source location.
    """
    lineno = getattr(node, "lineno", None)
    col_offset = getattr(node, "col_offset", None)
    end_lineno = getattr(node, "end_lineno", None)
    end_col_offset = getattr(node, "end_col_offset", None)
    if not (
        isinstance(lineno, int)
        and isinstance(col_offset, int)
        and isinstance(end_lineno, int)
        and isinstance(end_col_offset, int)
    ):
        raise NodeSpanError(f"{type(node).__name__} has no source location")
    return lineno, col_offset, end_lineno, end_col_offset


@traces(spec="110", story="US-110-009", reqs=["REQ-110-123"])
def _line_starts(source: bytes) -> tuple[int, ...]:
    """Byte offset of the first byte of each source line.

    Args:
        source: Raw file bytes.

    Returns:
        Tuple where ``result[lineno - 1]`` is the byte offset of the start
        of 1-indexed line ``lineno`` -- add an ast ``col_offset`` to get
        an absolute byte offset.
    """
    first = len(_UTF8_BOM) if source.startswith(_UTF8_BOM) else 0
    starts = [first]
    for match in _LINE_TERMINATOR_RE.finditer(source):
        starts.append(match.end())
    return tuple(starts)


@traces(spec="110", story="US-110-009", reqs=["REQ-110-123"])
@dataclass(frozen=True, slots=True)
class NodeIndex:
    """Maps ast nodes to their token byte-spans over a TokenStream.

    Attributes:
        stream: The TokenStream this index was built from.
        line_starts: Byte offset of each source line's first byte
            (``line_starts[lineno - 1]``).
        start_index: ``token.start -> index into stream.tokens``, LAST
            token at that offset wins. Zero-width structural tokens
            (INDENT/DEDENT) are emitted immediately BEFORE the real token
            they announce and share its start offset (e.g. a DEDENT right
            before a dedented ``def`` after a blank line) -- taking the
            last token at a start offset skips past them to the real one.
        end_index: ``token.end -> index into stream.tokens``, FIRST token
            ending there wins (the mirror case: a real token's end can
            coincide with a zero-width DEDENT/NEWLINE emitted right after
            it, and the real token -- emitted first -- is the one wanted).
    """

    stream: TokenStream
    line_starts: tuple[int, ...]
    start_index: dict[int, int]
    end_index: dict[int, int]

    @traces(spec="110", story="US-110-009", reqs=["REQ-110-123"])
    @staticmethod
    def build(tree: ast.Module, stream: TokenStream) -> NodeIndex:
        """Build a NodeIndex over ``stream`` for ``tree``.

        Args:
            tree: The ``ast.Module`` parsed from the SAME source as
                ``stream``. Not walked eagerly -- queries resolve each
                node's span lazily from its own location fields, so this
                is only accepted to keep the ast/token pairing explicit
                at the call site (mirrors ``LosslessModule.parse``).
            stream: The lossless TokenStream for that source.

        Returns:
            A NodeIndex ready for ``span_for``/``token_for_name`` queries.
        """
        del tree
        start_index: dict[int, int] = {}
        end_index: dict[int, int] = {}
        for i, token in enumerate(stream.tokens):
            start_index[token.start] = i  # last wins -- see start_index docstring
            end_index.setdefault(token.end, i)  # first wins -- see end_index docstring
        return NodeIndex(
            stream=stream,
            line_starts=_line_starts(stream.source),
            start_index=start_index,
            end_index=end_index,
        )

    @traces(spec="110", story="US-110-009", reqs=["REQ-110-123"])
    def _offset(self, lineno: int, col_offset: int) -> int:
        """Absolute byte offset for an ast ``(lineno, col_offset)`` pair."""
        return self.line_starts[lineno - 1] + col_offset

    @traces(spec="110", story="US-110-009", reqs=["REQ-110-123"])
    def span_for(self, node: ast.AST) -> Token:
        """The token span covering ``node``'s full source range.

        Args:
            node: Any ast node carrying source location (``lineno``/
                ``col_offset``/``end_lineno``/``end_col_offset``).

        Returns:
            A ``Token`` spanning ``node``'s bytes. ``.kind`` is the kind
            of the FIRST token in that span -- a representative tag, not
            a claim that the whole span is one lexical token (most ast
            nodes span many tokens).

        Raises:
            NodeSpanError: ``node`` has no location info, or its start
                byte offset does not align with any token boundary.
        """
        lineno, col_offset, end_lineno, end_col_offset = _location(node)
        start = self._offset(lineno, col_offset)
        end = self._offset(end_lineno, end_col_offset)
        index = self.start_index.get(start)
        if index is None:
            raise NodeSpanError(
                f"{type(node).__name__} span start {start} does not align with "
                "any token boundary"
            )
        kind = self.stream.tokens[index].kind
        return Token(kind=kind, start=start, end=end)

    def token_for_name(self, node: ast.AST) -> Token:
        """The exact NAME token that is ``node``'s rename target.

        Args:
            node: An ``ast.Name``, ``ast.Attribute``, ``ast.FunctionDef``,
                ``ast.AsyncFunctionDef``, or ``ast.ClassDef``.

        Returns:
            The single NAME token spanning the identifier itself: for
            ``Name``, the identifier token; for ``Attribute``, the
            trailing ``.attr`` name (NOT the value expression); for a
            def/class, the name after the ``def``/``class``/``async def``
            keyword(s) -- ast def-names have no span of their own, so
            this walks forward from the keyword and skips every leading
            keyword token via ``keyword.iskeyword`` (handles ``async
            def`` needing to skip TWO keyword tokens, not just one).

        Raises:
            NodeSpanError: the identifier token cannot be located.
            TypeError: ``node`` is not one of the supported node types.
        """
        if isinstance(node, ast.Name):
            lineno, col_offset, _, _ = _location(node)
            start = self._offset(lineno, col_offset)
            return self._name_token_at(self.start_index.get(start), node)
        if isinstance(node, ast.Attribute):
            _, _, end_lineno, end_col_offset = _location(node)
            end = self._offset(end_lineno, end_col_offset)
            return self._name_token_at(self.end_index.get(end), node)
        if isinstance(node, _KEYWORD_NAMED_TYPES):
            lineno, col_offset, _, _ = _location(node)
            start = self._offset(lineno, col_offset)
            index = self.start_index.get(start)
            if index is None:
                raise NodeSpanError(
                    f"{type(node).__name__} keyword start {start} does not align "
                    "with any token boundary"
                )
            for token in self.stream.tokens[index:]:
                if token.kind is TokenKind.NAME:
                    text = token.text(self.stream.source).decode("utf-8")
                    if not keyword.iskeyword(text):
                        return token
            raise NodeSpanError(
                f"no identifier NAME token found after {type(node).__name__} "
                "keyword(s)"
            )
        raise TypeError(f"token_for_name does not support {type(node).__name__}")

    def _name_token_at(self, index: int | None, node: ast.AST) -> Token:
        """Resolve a token index to a NAME token, or raise NodeSpanError."""
        if index is None:
            raise NodeSpanError(f"{type(node).__name__} has no aligned token")
        token = self.stream.tokens[index]
        if token.kind is not TokenKind.NAME:
            raise NodeSpanError(
                f"{type(node).__name__} resolved to a {token.kind} token, not NAME"
            )
        return token