Add targeted source map lookup

README.md

@@ -54,6 +54,7 @@ positional information (`at line N, column M`) and re-raises them in the same
 
 ```python
 from parse_errors import ParseContext, ParseError
+from parse_errors.source_map import SourceMap, build_source_map, locate_pointer
 ```
 
 **`ParseContext(filename, *, data=None, format=None)`** — context manager.
@@ -68,6 +69,28 @@ from parse_errors import ParseContext, ParseError
 **`ParseError`** — the exception raised inside the context.  Has attributes
 `filename`, `line` (1-based), and `column` (1-based).
 
+**`locate_pointer(source, fmt, pointer)`** — returns the best source-map entry
+for one JSON Pointer without building a full map.  `fmt` is `"json"`,
+`"toml"`, or `"yaml"`; `source` is `str` or UTF-8 `bytes`; `pointer` uses RFC
+6901 escaping.  If the exact pointer is not present, the result matches
+`closest_entry(build_source_map(source, fmt), pointer)`: the nearest enclosing
+value when one exists, otherwise `None`.
+
+**`SourceMap(source, fmt)`** — caches the parsed document for repeated targeted
+lookups.  Use `SourceMap(...).locate(pointer)` when several errors in the same
+document need locations.  It still avoids constructing a full pointer-to-entry
+map.
+
+**`build_source_map(source, fmt)`** — builds the full pointer-to-location map.
+This is useful when callers need many arbitrary entries or need to inspect all
+locations.
+
+**Warning:** source-map helpers locate nodes in a document; they are not
+validating parsers.  JSON and TOML location support uses tree-sitter so it can
+return a location from a syntax tree even when a real decoder would reject the
+source.  Parse or validate the document with your normal parser first, then use
+these helpers only to map known error paths back to source locations.
+
 # Version Compat
 
 This library is compatile with Python 3.10+, but should be linted under the

parse_errors/context.py

@@ -8,7 +8,7 @@
 from pathlib import Path
 from typing import Iterator
 
-from .source_map import detect_format, build_source_map, closest_entry, Location
+from .source_map import Location, detect_format, locate_pointer
 from ._jsonpath import extract_jsonpath, jsonpath_to_pointer
 
 POSITIONAL_RE = re.compile(r"at line (\d+), column (\d+)")
@@ -91,9 +91,7 @@ def ParseContext(
         assert fmt is not None
 
         source = data if data is not None else path.read_bytes()
-        source_map = build_source_map(source, fmt)
-
-        entry = closest_entry(source_map, pointer)
+        entry = locate_pointer(source, fmt, pointer)
         if entry is None:  # pragma: no cover
             raise exc
 

parse_errors/json_source_map/__init__.py

@@ -1,25 +1,222 @@
-"""Thin wrapper around the json-source-map package that returns our own types."""
+"""Calculate JSON source locations using tree-sitter, targeted or full.
+
+tree-sitter does the parse in C instead of a hand-written pure-Python
+character scanner, and locate() additionally never visits a subtree that
+isn't on the path to the one pointer asked for -- indexing directly into
+tree-sitter's parsed node tree.
+
+calculate() (the full map, used by anything that still wants build_source_map
++ closest_entry) is kept for comparison/completeness.
+"""
 
 from __future__ import annotations
 
-import json_source_map as _ext
+import json
+
+import tree_sitter_json
+import tree_sitter as ts
 
 from ..source_map import Entry, Location, TSourceMap
+from .._jsonpath import _escape
+
+_LANGUAGE = ts.Language(tree_sitter_json.language())
+
+_Parsed = tuple[ts.Node, list[bytes]]
+
+
+def calculate(source: str | bytes) -> TSourceMap:
+    """Calculate the full source map for a JSON document."""
+    src_bytes = source.encode("utf-8") if isinstance(source, str) else source
+    root = ts.Parser(_LANGUAGE).parse(src_bytes).root_node
+    lines = src_bytes.splitlines(True) + [b""]
+    byte_to_char = [_byte_to_char_offsets(line) for line in lines]
+    line_start_chars = [0]
+    for offsets in byte_to_char[:-1]:
+        line_start_chars.append(line_start_chars[-1] + offsets[-1])
+
+    result: TSourceMap = {}
+    top = root.named_children[0] if root.named_children else root
+    _walk_all(top, [], result, byte_to_char, line_start_chars)
+    return result
+
+
+def _parse(source: str | bytes) -> _Parsed | None:
+    """Parse *source* into the (node, lines) pair :func:`_locate_in` walks.
+
+    Split out from :func:`locate` so :class:`~parse_errors.source_map.SourceMap`
+    can parse once and reuse it across several ``locate()`` calls on the same
+    document, instead of reparsing per pointer.
+    """
+    src_bytes = source.encode("utf-8") if isinstance(source, str) else source
+    root = ts.Parser(_LANGUAGE).parse(src_bytes).root_node
+    lines = src_bytes.splitlines(True) + [b""]
+    top = root.named_children[0] if root.named_children else root
+    return top, lines
+
+
+def _locate_in(parsed: _Parsed | None, pointer: str) -> Entry | None:
+    if parsed is None:
+        return None
+    top, lines = parsed
+    segments = _pointer_segments(pointer)
+    return _walk_targeted(top, segments, 0, lines)
+
+
+def locate(source: str | bytes, pointer: str) -> Entry | None:
+    """Find *pointer*'s location, visiting only nodes on the path to it."""
+    return _locate_in(_parse(source), pointer)
+
+
+# --- shared helpers ---
+
 
+def _pointer_segments(pointer: str) -> list[str]:
+    if not pointer:
+        return []
+    return [p.replace("~1", "/").replace("~0", "~") for p in pointer.split("/")[1:]]
 
-def calculate(source: str) -> TSourceMap:
-    """Calculate the source map for a JSON document."""
-    return {
-        pointer: Entry(
-            value_start=_loc(e.value_start),
-            value_end=_loc(e.value_end),
-            key_start=_loc(e.key_start) if e.key_start is not None else None,
-            key_end=_loc(e.key_end) if e.key_end is not None else None,
+
+def _to_pointer(segments: list[str]) -> str:
+    return "/" + "/".join(_escape(s) for s in segments) if segments else ""
+
+
+def _byte_to_char_offsets(line: bytes) -> list[int]:
+    offsets = [0] * (len(line) + 1)
+    char_count = 0
+    for i, b in enumerate(line):
+        if b < 0x80 or b >= 0xC0:
+            char_count += 1
+        offsets[i + 1] = char_count
+    return offsets
+
+
+# --- full map (calculate) ---
+
+
+def _walk_all(
+    node: ts.Node,
+    segments: list[str],
+    result: TSourceMap,
+    byte_to_char: list[list[int]],
+    line_start_chars: list[int],
+    key_node: ts.Node | None = None,
+) -> None:
+    pointer = _to_pointer(segments)
+    result[pointer] = Entry(
+        value_start=_loc_precomputed(node.start_point, byte_to_char, line_start_chars),
+        value_end=_loc_precomputed(node.end_point, byte_to_char, line_start_chars),
+        key_start=_loc_precomputed(key_node.start_point, byte_to_char, line_start_chars)
+        if key_node
+        else None,
+        key_end=_loc_precomputed(key_node.end_point, byte_to_char, line_start_chars)
+        if key_node
+        else None,
+    )
+    if node.type == "object":
+        for pair in node.named_children:
+            if pair.type != "pair":
+                continue  # pragma: no cover
+            k = pair.child_by_field_name("key")
+            v = pair.child_by_field_name("value")
+            assert k is not None and k.text is not None
+            assert v is not None
+            key_text = json.loads(k.text.decode("utf-8"))
+            _walk_all(
+                v,
+                segments + [key_text],
+                result,
+                byte_to_char,
+                line_start_chars,
+                key_node=k,
+            )
+    elif node.type == "array":
+        for i, child in enumerate(node.named_children):
+            _walk_all(
+                child, segments + [str(i)], result, byte_to_char, line_start_chars
+            )
+
+
+def _loc_precomputed(
+    point: ts.Point, byte_to_char: list[list[int]], line_start_chars: list[int]
+) -> Location:
+    char_column = byte_to_char[point.row][point.column]
+    return Location(
+        line=point.row,
+        column=char_column,
+        position=line_start_chars[point.row] + char_column,
+    )
+
+
+# --- targeted lookup ---
+
+
+def _walk_targeted(
+    node: ts.Node,
+    segments: list[str],
+    idx: int,
+    lines: list[bytes],
+    key_node: ts.Node | None = None,
+) -> Entry:
+    if idx == len(segments):
+        return Entry(
+            value_start=_loc(node.start_point, lines),
+            value_end=_loc(node.end_point, lines),
+            key_start=_loc(key_node.start_point, lines)
+            if key_node is not None
+            else None,
+            key_end=_loc(key_node.end_point, lines) if key_node is not None else None,
         )
-        for pointer, e in _ext.calculate(source).items()
-    }
+
+    if node.type == "object":
+        target_key = segments[idx]
+        for pair in node.named_children:
+            if pair.type != "pair":
+                continue  # pragma: no cover
+            k = pair.child_by_field_name("key")
+            assert k is not None and k.text is not None
+            key_text = json.loads(k.text.decode("utf-8"))
+            if key_text == target_key:
+                v = pair.child_by_field_name("value")
+                assert v is not None
+                return _walk_targeted(v, segments, idx + 1, lines, key_node=k)
+        # target_key not present -- this object is the closest ancestor
+        return _ancestor(node, lines, key_node)
+
+    if node.type == "array":
+        try:
+            target_index = int(segments[idx])
+        except ValueError:
+            target_index = -1
+        children = node.named_children
+        if 0 <= target_index < len(children):
+            return _walk_targeted(children[target_index], segments, idx + 1, lines)
+        return _ancestor(node, lines, key_node)
+
+    # A scalar, but the pointer wants to go deeper -- no such path; this
+    # value is the closest ancestor.
+    return _ancestor(node, lines, key_node)
+
+
+def _ancestor(node: ts.Node, lines: list[bytes], key_node: ts.Node | None) -> Entry:
+    # A fallback entry still carries key_start/key_end when this node is
+    # itself the value of an enclosing key -- matching the entry it would
+    # have gotten in the full map, at the pointer for that key.
+    return Entry(
+        value_start=_loc(node.start_point, lines),
+        value_end=_loc(node.end_point, lines),
+        key_start=_loc(key_node.start_point, lines) if key_node is not None else None,
+        key_end=_loc(key_node.end_point, lines) if key_node is not None else None,
+    )
 
 
-def _loc(ext: _ext.Location) -> Location:
-    """Translate to our internal structure."""
-    return Location(line=ext.line, column=ext.column, position=ext.position)
+def _loc(point: ts.Point, lines: list[bytes]) -> Location:
+    # Only called O(depth) times per locate() call (a handful), so decoding
+    # fresh each time -- rather than precomputing a whole-document table --
+    # is cheap; the table only paid for itself when amortized over every
+    # node in the document, which locate() deliberately never visits.
+    line = lines[point.row]
+    char_column = len(line[: point.column].decode("utf-8"))
+    position = (
+        sum(len(prior.decode("utf-8")) for prior in lines[: point.row]) + char_column
+    )
+    return Location(line=point.row, column=char_column, position=position)

parse_errors/source_map.py

@@ -4,7 +4,7 @@
 
 import dataclasses
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Any, Dict, Mapping, Optional, TypeVar
 
 
 # These are identical to the ones in json-source-map, but I feel icky exporting
@@ -27,6 +27,7 @@ class Entry:
 
 
 TSourceMap = Dict[str, Entry]
+TEntry = TypeVar("TEntry")
 
 
 def detect_format(path: Path) -> str | None:
@@ -62,7 +63,85 @@ def build_source_map(source: str | bytes, fmt: str) -> TSourceMap:
         raise ValueError(f"Unknown format: {fmt!r}")
 
 
-def closest_entry(source_map: TSourceMap, pointer: str) -> Entry | None:
+def _parse_for(fmt: str, source: str | bytes) -> Any:
+    # The three formats' parse trees (tree-sitter nodes for json/toml, a
+    # PyYAML Node for yaml) share no common type -- this is deliberately
+    # opaque, threaded straight to the matching format's own _locate_in().
+    if fmt == "json":
+        from . import json_source_map
+
+        return json_source_map._parse(source)
+    if fmt == "toml":
+        from . import toml_source_map
+
+        return toml_source_map._parse(source)
+    if fmt in ("yaml", "yml"):
+        from . import yaml_source_map
+
+        return yaml_source_map._parse(source)
+    raise ValueError(f"Unknown format: {fmt!r}")
+
+
+def _locate_in_for(fmt: str, parsed: Any, pointer: str) -> Entry | None:
+    if fmt == "json":
+        from . import json_source_map
+
+        return json_source_map._locate_in(parsed, pointer)
+    if fmt == "toml":
+        from . import toml_source_map
+
+        return toml_source_map._locate_in(parsed, pointer)
+    if fmt in ("yaml", "yml"):
+        from . import yaml_source_map
+
+        return yaml_source_map._locate_in(parsed, pointer)
+    raise ValueError(f"Unknown format: {fmt!r}")
+
+
+def locate_pointer(source: str | bytes, fmt: str, pointer: str) -> Entry | None:
+    """Find the location of one pointer directly, without building the whole map.
+
+    A targeted walk over each format's own parse tree -- tree-sitter for
+    json/toml, PyYAML's Node tree (via ``get_single_node()``, CSafeLoader when
+    available) for yaml -- visiting only nodes on the path to *pointer*,
+    unlike :func:`build_source_map` which computes an entry for every value
+    in the document regardless of which one is wanted.
+
+    Looking up several pointers in the same document this way reparses it
+    each time; use :class:`SourceMap` instead to parse once and reuse it.
+    """
+    return _locate_in_for(fmt, _parse_for(fmt, source), pointer)
+
+
+class SourceMap:
+    """Repeated pointer lookups against one document, parsed at most once.
+
+    :func:`locate_pointer` costs O(pointer depth) per call, but reparses the
+    document from scratch every time -- fine for one lookup, wasteful for
+    several against the same source, such as a merged config reporting one
+    error per field. ``SourceMap`` defers parsing until the first
+    :meth:`locate` call, caches that parse, and reuses it for every pointer
+    asked after -- still never building a whole-document map the way
+    :func:`build_source_map` does, since each lookup still only walks the
+    path to the one pointer it was asked for.
+    """
+
+    __slots__ = ("_source", "_fmt", "_parsed", "_has_parsed")
+
+    def __init__(self, source: str | bytes, fmt: str) -> None:
+        self._source = source
+        self._fmt = fmt
+        self._parsed: Any = None
+        self._has_parsed = False
+
+    def locate(self, pointer: str) -> Entry | None:
+        if not self._has_parsed:
+            self._parsed = _parse_for(self._fmt, self._source)
+            self._has_parsed = True
+        return _locate_in_for(self._fmt, self._parsed, pointer)
+
+
+def closest_entry(source_map: Mapping[str, TEntry], pointer: str) -> TEntry | None:
     """Return the source map entry for ``pointer``, falling back to the longest prefix."""
     if pointer in source_map:
         return source_map[pointer]