scripts/tree_utils.py

"""Structural parsing shared by data generation, training, and evaluation."""

from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass

REF_RE = re.compile(r"\[ref=([^\]]+)\]")
STAGEHAND_REF_RE = re.compile(r"^\s*\[([^\]]+)\]")
WORD_RE = re.compile(r"[a-z0-9]{2,}")


@dataclass(frozen=True)
class Window:
    id: str
    text: str
    refs: tuple[str, ...]
    start_line: int
    end_line: int


def references(text: str, source_format: str = "betterwright") -> set[str]:
    if source_format == "stagehand":
        return {m.group(1) for line in text.splitlines() if (m := STAGEHAND_REF_RE.match(line))}
    # Auto-detect mixed corpora: Stagehand's outline is `[id] role: name`,
    # while BetterWright/Playwright use `[ref=id]`. The patterns cannot collide.
    result = set(REF_RE.findall(text))
    result.update(
        m.group(1) for line in text.splitlines() if (m := STAGEHAND_REF_RE.match(line))
    )
    return result


def words(text: str) -> set[str]:
    return set(WORD_RE.findall(text.casefold()))


def structural_windows(text: str, *, max_chars: int = 3600, overlap_lines: int = 4) -> list[Window]:
    """Create overlapping, indentation-aware windows without cutting a line."""
    lines = [line.rstrip() for line in text.splitlines() if line.strip()]
    if not lines:
        return []
    windows: list[Window] = []
    start = 0
    while start < len(lines):
        size = 0
        end = start
        while end < len(lines) and (size + len(lines[end]) + 1 <= max_chars or end == start):
            size += len(lines[end]) + 1
            end += 1
        # Prefer a structural boundary near the limit.
        if end < len(lines):
            floor = max(start + 1, end - 12)
            candidates = [i for i in range(floor, end) if len(lines[i]) - len(lines[i].lstrip()) <= 2]
            if candidates:
                end = candidates[-1]
        body = "\n".join(lines[start:end])
        refs = tuple(sorted(references(body)))
        digest = hashlib.sha1(f"{start}\0{end}\0{body}".encode()).hexdigest()[:16]
        windows.append(Window(digest, body, refs, start, end))
        if end >= len(lines):
            break
        start = max(start + 1, end - overlap_lines)
    return windows


def lexical_score(query: str, text: str) -> float:
    q = words(query)
    if not q:
        return 0.0
    t = words(text)
    overlap = len(q & t)
    return overlap / (len(q) ** 0.5 * max(1, len(t)) ** 0.25)