codebase-receipts-cli 1.0.0__py3-none-any.whl

receipts/ingest/artifact_extractor.py

@@ -0,0 +1,679 @@
+"""tree-sitter based extraction of key code artifacts within a token budget.
+
+API facts verified against the installed tree-sitter 0.26.0 and
+tree-sitter-language-pack 1.12.2: parsers come from
+``tree_sitter.Parser(get_language(name))`` (the pack's own ``get_parser``
+returns an incompatible native class), ``parse()`` requires bytes, and
+decorated Python defs appear as ``decorated_definition`` wrapping the actual
+``function_definition``/``class_definition``.
+
+Every artifact's code is redacted with the secrets scanner before it leaves
+this module — content headed for the knowledge base must never carry a
+credential. Budget overruns are counted and reported, never silent.
+"""
+
+from __future__ import annotations
+
+import re
+from collections import deque
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+
+try:
+    import tree_sitter
+    from tree_sitter_language_pack import get_language
+except (ImportError, OSError) as _tree_sitter_err:
+    import platform as _platform
+    import sys as _sys
+
+    _arch = _platform.machine()
+    _os = _platform.system()
+    _py = _sys.version.split()[0]
+    raise SystemExit(
+        f"\n[receipts] tree-sitter failed to load on {_os}/{_arch} (Python {_py}).\n"
+        "\n"
+        "This usually means a prebuilt wheel isn't available for your\n"
+        "architecture. To fix it you need a C compiler so pip can build\n"
+        "the extension from source:\n"
+        "\n"
+        "  Linux/ARM64 (e.g. Raspberry Pi, Graviton):\n"
+        "    sudo apt install build-essential python3-dev\n"
+        "    pip install tree-sitter tree-sitter-language-pack\n"
+        "\n"
+        "  macOS:\n"
+        "    xcode-select --install\n"
+        "    pip install tree-sitter tree-sitter-language-pack\n"
+        "\n"
+        "  Windows (no C compiler):\n"
+        "    Install Visual Studio Build Tools from:\n"
+        "    https://visualstudio.microsoft.com/visual-cpp-build-tools/\n"
+        "    then re-run: pip install tree-sitter tree-sitter-language-pack\n"
+        "\n"
+        "If your platform is genuinely unsupported, open an issue at:\n"
+        "https://github.com/Ishaan-1606/receipts-cli/issues\n"
+        "\n"
+        f"Original error: {_tree_sitter_err}\n"
+    ) from None
+
+from receipts.ingest import scanner
+from receipts.ingest.secrets_scanner import redact
+
+_EXTENSION_LANGUAGES = {
+    ".py": "python",
+    ".js": "javascript",
+    ".jsx": "javascript",
+    ".ts": "typescript",
+    ".tsx": "tsx",
+    # Phase 20 — extraction parity across mainstream languages, chosen by
+    # usage share. Every language here has BOTH a tree-sitter map (verified
+    # against the installed tree-sitter-language-pack by probing each
+    # grammar's actual node types) and a regex signature fallback (the path
+    # Windows uses, where tree-sitter is disabled).
+    ".java": "java",
+    ".kt": "kotlin",
+    ".kts": "kotlin",
+    ".go": "go",
+    ".rs": "rust",
+    ".c": "c",
+    ".h": "c",
+    ".cpp": "cpp",
+    ".cc": "cpp",
+    ".cxx": "cpp",
+    ".hpp": "cpp",
+    ".hh": "cpp",
+    ".cs": "csharp",
+    ".php": "php",
+    ".rb": "ruby",
+    ".swift": "swift",
+}
+
+_PY_DEF_TYPES = {"function_definition": "function", "class_definition": "class"}
+_JS_DEF_TYPES = {
+    "function_declaration": "function",
+    "class_declaration": "class",
+    "interface_declaration": "interface",
+}
+
+# Top-level definition node types per grammar → artifact kind. Every entry
+# was verified by parsing real snippets with the installed language pack —
+# never assumed from memory (several grammars surprised: Kotlin has no
+# `name` field, Go type names live on the inner type_spec, Swift parses
+# extensions as class_declaration, C#/C++ nest types inside namespaces).
+_TS_DEF_TYPES: dict[str, dict[str, str]] = {
+    "python": _PY_DEF_TYPES,
+    "javascript": _JS_DEF_TYPES,
+    "typescript": _JS_DEF_TYPES,
+    "tsx": _JS_DEF_TYPES,
+    "java": {
+        "class_declaration": "class",
+        "interface_declaration": "interface",
+        "enum_declaration": "class",
+        "record_declaration": "class",
+    },
+    "kotlin": {
+        "function_declaration": "function",
+        "class_declaration": "class",
+        "object_declaration": "class",
+    },
+    "go": {
+        "function_declaration": "function",
+        "method_declaration": "function",
+        "type_declaration": "class",
+    },
+    "rust": {
+        "function_item": "function",
+        "struct_item": "class",
+        "enum_item": "class",
+        "trait_item": "interface",
+        "impl_item": "class",
+    },
+    "c": {
+        "function_definition": "function",
+        "struct_specifier": "class",
+        "enum_specifier": "class",
+    },
+    "cpp": {
+        "function_definition": "function",
+        "class_specifier": "class",
+        "struct_specifier": "class",
+        "namespace_definition": "class",
+        "enum_specifier": "class",
+    },
+    "csharp": {
+        # With block namespaces the whole namespace becomes one artifact
+        # (like a Python class with its methods); file-scoped namespaces
+        # leave types at top level where they match individually.
+        "namespace_declaration": "class",
+        "class_declaration": "class",
+        "interface_declaration": "interface",
+        "struct_declaration": "class",
+        "enum_declaration": "class",
+        "record_declaration": "class",
+    },
+    "php": {
+        "class_declaration": "class",
+        "interface_declaration": "interface",
+        "trait_declaration": "class",
+        "function_definition": "function",
+    },
+    "ruby": {
+        "class": "class",
+        "module": "class",
+        "method": "function",
+    },
+    "swift": {
+        "class_declaration": "class",  # covers struct/enum/extension too
+        "protocol_declaration": "interface",
+        "function_declaration": "function",
+    },
+}
+
+#: Fallback excerpt size for files with no extractable definitions
+#: (README, config, unsupported languages).
+_FALLBACK_MAX_LINES = 150
+
+#: Lines that are import/package boilerplate — the fallback excerpt skips a
+#: contiguous run of these (plus blanks) at the top of the file so the
+#: excerpt captures logic instead of headers. Deliberately narrow: comment
+#: markers are NOT skipped (a markdown title or license header may be the
+#: most informative content a doc file has).
+_BOILERPLATE_LINE_RE = re.compile(
+    r"^\s*(?:import\b|from\s+[\w.]+\s+import\b|package\b|using\s+\w|"
+    r"#include\b|require\s*[('\"]|require_relative\b|use\s+[A-Za-z\\]|"
+    r"<\?php\s*$|\s*$)"
+)
+
+#: Ceiling for the auto-scaled budget (token_budget=None). Corpora under it
+#: are ingested whole; beyond it, importance allocation decides what stays.
+_AUTO_BUDGET_CEILING = 1_000_000
+
+
+@dataclass(frozen=True)
+class Artifact:
+    rel_path: str
+    language: str  # "python", "javascript", ... or "text" for fallbacks
+    kind: str  # "function" | "class" | "interface" | "file"
+    name: str
+    start_line: int  # 1-based
+    end_line: int
+    code: str
+    token_estimate: int
+
+
+@dataclass
+class ExtractionResult:
+    budget: int
+    artifacts: list[Artifact] = field(default_factory=list)
+    used_tokens: int = 0
+    dropped_over_budget: int = 0
+    redaction_count: int = 0
+    auto_budget: bool = False  # True when the budget was auto-scaled
+    corpus_tokens: int = 0  # total tokens across ALL candidates, pre-budget
+    dropped_tokens: int = 0
+    dropped_by_dir: dict[str, int] = field(default_factory=dict)
+
+
+_parsers: dict[str, tree_sitter.Parser] = {}
+
+
+def _parser_for(language: str) -> tree_sitter.Parser:
+    if language not in _parsers:
+        _parsers[language] = tree_sitter.Parser(get_language(language))
+    return _parsers[language]
+
+
+#: Node types that ARE the identifier, for grammars without a `name` field
+#: (verified: kotlin uses simple_identifier/type_identifier, rust impl uses
+#: type_identifier, go type_declaration wraps a type_spec, ruby uses
+#: constant/identifier).
+_IDENTIFIER_TYPES = {
+    "identifier",
+    "simple_identifier",
+    "type_identifier",
+    "field_identifier",
+    "namespace_identifier",
+    "constant",
+    "name",
+}
+
+
+def _node_name(node: tree_sitter.Node, depth: int = 0) -> str:
+    name_node = node.child_by_field_name("name")
+    if name_node is not None and name_node.text is not None:
+        return name_node.text.decode("utf-8", errors="replace")
+    # Grammars without a `name` field: take the first identifier-ish child,
+    # recursing a little for wrapper nodes (go type_spec, c declarators).
+    for child in node.named_children:
+        if child.type in _IDENTIFIER_TYPES and child.text is not None:
+            return child.text.decode("utf-8", errors="replace")
+    if depth < 2:
+        for child in node.named_children:
+            found = _node_name(child, depth + 1)
+            if found != "<anonymous>":
+                return found
+    return "<anonymous>"
+
+
+#: Namespace containers whose members should be extracted individually —
+#: one verified level of recursion into their declaration_list, so a C#
+#: block namespace yields its classes instead of one namespace-sized blob.
+_NAMESPACE_TYPES = {
+    "csharp": "namespace_declaration",
+    "cpp": "namespace_definition",
+}
+
+
+def _definition_nodes(
+    root: tree_sitter.Node, language: str
+) -> list[tuple[str, str, tree_sitter.Node]]:
+    """Yield (kind, name, node) for top-level defs; node spans decorators/export."""
+    defs: list[tuple[str, str, tree_sitter.Node]] = []
+    kind_map = _TS_DEF_TYPES.get(language, {})
+    for outer in root.named_children:
+        inner = outer
+        if language == "python" and outer.type == "decorated_definition":
+            definition = outer.child_by_field_name("definition")
+            if definition is not None:
+                inner = definition
+        elif (
+            language in ("javascript", "typescript", "tsx")
+            and outer.type == "export_statement"
+        ):
+            wrapped = next(
+                (c for c in outer.named_children if c.type in kind_map), None
+            )
+            if wrapped is not None:
+                inner = wrapped
+        elif outer.type == _NAMESPACE_TYPES.get(language):
+            body = next(
+                (c for c in outer.named_children if c.type == "declaration_list"),
+                None,
+            )
+            members = (
+                [
+                    (kind_map[c.type], _node_name(c), c)
+                    for c in body.named_children
+                    if c.type in kind_map
+                ]
+                if body is not None
+                else []
+            )
+            if members:
+                defs.extend(members)
+                continue
+            # Empty or unrecognized body — keep the namespace itself.
+        kind = kind_map.get(inner.type)
+        if kind is not None:
+            defs.append((kind, _node_name(inner), outer))
+    return defs
+
+
+def _ts_code_artifacts(
+    rel_path: str, language: str, text: str
+) -> list[tuple[str, str, int, int, str]]:
+    """tree-sitter extraction — may crash on Windows with large batches."""
+    source = text.encode("utf-8")
+    tree = _parser_for(language).parse(source)
+    results = []
+    for kind, name, node in _definition_nodes(tree.root_node, language):
+        code = source[node.start_byte : node.end_byte].decode("utf-8", errors="replace")
+        results.append(
+            (kind, name, node.start_point.row + 1, node.end_point.row + 1, code)
+        )
+    return results
+
+
+_PY_DEF_RE = re.compile(
+    r"^(?P<deco>(?:@\w[\w.]*(?:\(.*?\))?\s*\n)*)?"
+    r"^(?P<kw>(?:async\s+)?def|class)\s+(?P<name>\w+)",
+    re.MULTILINE,
+)
+_JS_DEF_RE = re.compile(
+    r"^(?:export\s+(?:default\s+)?)?(?P<kw>function|class|interface)"
+    r"\s+(?P<name>\w+)",
+    re.MULTILINE,
+)
+
+_KIND_MAP = {
+    "def": "function",
+    "async def": "function",
+    "class": "class",
+    "function": "function",
+    "interface": "interface",
+    "fun": "function",
+    "fn": "function",
+    "func": "function",
+    "object": "class",
+    "record": "class",
+    "enum": "class",
+    "struct": "class",
+    "union": "class",
+    "trait": "class",
+    "impl": "class",
+    "type": "class",
+    "module": "class",
+    "protocol": "interface",
+    "namespace": "class",
+}
+
+# Regex signature fallback per language: (pattern, default_kind) pairs. The
+# type-level keywords allow a little indentation because C#/Java/PHP types
+# usually sit inside a namespace/class body; the slicer below cuts from one
+# match to the next, so a class artifact naturally spans its methods.
+_TYPE_MODS = (
+    r"(?:(?:public|private|protected|internal|open|abstract|sealed|static"
+    r"|final|partial|data|inline|suspend|operator|export|readonly)\s+)*"
+)
+
+_REGEX_DEF_PATTERNS: dict[str, list[tuple[re.Pattern, str]]] = {
+    "python": [(_PY_DEF_RE, "function")],
+    "javascript": [(_JS_DEF_RE, "function")],
+    "typescript": [(_JS_DEF_RE, "function")],
+    "tsx": [(_JS_DEF_RE, "function")],
+    "java": [
+        (
+            re.compile(
+                r"^(?:@\w+(?:\([^)]*\))?[ \t]*\n)*[ \t]{0,8}"
+                + _TYPE_MODS
+                + r"(?P<kw>class|interface|enum|record)\s+(?P<name>\w+)",
+                re.MULTILINE,
+            ),
+            "class",
+        ),
+    ],
+    "kotlin": [
+        (
+            re.compile(
+                r"^[ \t]{0,8}(?:@\w+[ \t]*)*"
+                + _TYPE_MODS
+                + r"(?P<kw>fun|class|object|interface)\s+(?P<name>`?\w+`?)",
+                re.MULTILINE,
+            ),
+            "function",
+        ),
+    ],
+    "go": [
+        (
+            re.compile(r"^func\s+(?:\([^)]*\)\s*)?(?P<name>\w+)", re.MULTILINE),
+            "function",
+        ),
+        (
+            re.compile(r"^type\s+(?P<name>\w+)\s+(?:struct|interface)\b", re.MULTILINE),
+            "class",
+        ),
+    ],
+    "rust": [
+        (
+            re.compile(
+                r"^[ \t]{0,4}(?:pub(?:\([^)]*\))?\s+)?"
+                r"(?:(?:async|unsafe|const|extern)\s+)*"
+                r"(?P<kw>fn|struct|enum|trait|impl)\s+(?P<name>\w+)",
+                re.MULTILINE,
+            ),
+            "function",
+        ),
+    ],
+    "c": [
+        (
+            re.compile(
+                r"^(?P<kw>struct|enum|union)\s+(?P<name>\w+)\s*\{", re.MULTILINE
+            ),
+            "class",
+        ),
+        (
+            # Column-0 "returntype name(" — top-level function definitions.
+            r"^(?!\s)(?:[A-Za-z_][\w \t\*]*[ \t\*])(?P<name>[A-Za-z_]\w*)\s*\(",
+            "function",
+        ),
+    ],
+    "cpp": [
+        (
+            re.compile(
+                r"^[ \t]{0,8}(?:template\s*<[^>\n]*>[ \t]*\n)?"
+                r"(?P<kw>class|struct|namespace)\s+(?P<name>\w+)",
+                re.MULTILINE,
+            ),
+            "class",
+        ),
+        (
+            r"^(?!\s)(?:[A-Za-z_][\w:<> \t\*&]*[ \t\*&])(?P<name>[A-Za-z_]\w*)\s*\(",
+            "function",
+        ),
+    ],
+    "csharp": [
+        (
+            re.compile(
+                r"^[ \t]{0,8}(?:\[[^\]\n]+\][ \t]*\n)*[ \t]{0,8}"
+                + _TYPE_MODS
+                + r"(?P<kw>class|interface|struct|enum|record)\s+(?P<name>\w+)",
+                re.MULTILINE,
+            ),
+            "class",
+        ),
+    ],
+    "php": [
+        (
+            re.compile(
+                r"^[ \t]{0,4}(?:(?:final|abstract)\s+)?"
+                r"(?P<kw>class|interface|trait|function)\s+(?P<name>\w+)",
+                re.MULTILINE,
+            ),
+            "function",
+        ),
+    ],
+    "ruby": [
+        (
+            re.compile(
+                r"^[ \t]{0,8}(?P<kw>class|module|def)\s+(?P<name>[\w.]+[?!]?)",
+                re.MULTILINE,
+            ),
+            "function",
+        ),
+    ],
+    "swift": [
+        (
+            re.compile(
+                r"^[ \t]{0,8}"
+                + _TYPE_MODS
+                + r"(?P<kw>func|class|struct|enum|protocol|extension)"
+                + r"\s+(?P<name>\w+)",
+                re.MULTILINE,
+            ),
+            "function",
+        ),
+    ],
+}
+
+# Compile the string patterns left raw above (kept raw for line length).
+_REGEX_DEF_PATTERNS = {
+    lang: [
+        (re.compile(p, re.MULTILINE) if isinstance(p, str) else p, kind)
+        for p, kind in specs
+    ]
+    for lang, specs in _REGEX_DEF_PATTERNS.items()
+}
+
+
+def _regex_code_artifacts(
+    rel_path: str, language: str, text: str
+) -> list[tuple[str, str, int, int, str]]:
+    """Regex-based fallback that never crashes. Less precise than tree-sitter
+    but catches top-level definitions reliably — and it's the ONLY code path
+    on Windows, where tree-sitter is disabled."""
+    specs = _REGEX_DEF_PATTERNS.get(language)
+    if not specs:
+        return []
+
+    found: dict[int, tuple[str, str]] = {}  # start_pos -> (kind, name)
+    for pattern, default_kind in specs:
+        for m in pattern.finditer(text):
+            kw = (m.groupdict().get("kw") or "").strip()
+            kind = _KIND_MAP.get(kw, default_kind)
+            found.setdefault(m.start(), (kind, m.group("name")))
+
+    if not found:
+        return []
+
+    starts = sorted(found)
+    results = []
+    for i, start in enumerate(starts):
+        kind, name = found[start]
+        end_pos = starts[i + 1] if i + 1 < len(starts) else len(text)
+        start_line = text[:start].count("\n") + 1
+        end_line = text[:end_pos].count("\n") + 1
+        code = text[start:end_pos].rstrip()
+        results.append((kind, name, start_line, end_line, code))
+    return results
+
+
+import platform as _platform  # noqa: E402  (deliberately after the fallback defs)
+
+_USE_TREE_SITTER = _platform.system() != "Windows"
+
+
+def _code_artifacts(
+    rel_path: str, language: str, text: str
+) -> list[tuple[str, str, int, int, str]]:
+    """Extract code artifacts. Uses tree-sitter on Linux/macOS, regex on
+    Windows (tree-sitter-language-pack has native crashes on Windows when
+    parsing many files sequentially)."""
+    if _USE_TREE_SITTER:
+        try:
+            return _ts_code_artifacts(rel_path, language, text)
+        except Exception:
+            pass
+    return _regex_code_artifacts(rel_path, language, text)
+
+
+def _fallback_artifact(rel_path: str, text: str) -> tuple[str, str, int, int, str]:
+    lines = text.splitlines()
+
+    # Skip the contiguous import/package block at the file head — in most
+    # compiled languages the first 30 lines are boilerplate, and an excerpt
+    # of boilerplate produces "Unsupported" verdicts for the wrong reason.
+    start = 0
+    for i, line in enumerate(lines[:80]):
+        if not _BOILERPLATE_LINE_RE.match(line):
+            start = i
+            break
+    else:
+        start = 0 if len(lines) <= 80 else 80
+
+    end = min(len(lines), start + _FALLBACK_MAX_LINES)
+    excerpt = "\n".join(lines[start:end])
+    name = rel_path.rsplit("/", 1)[-1]
+    return ("file", name, start + 1, max(end, 1), excerpt)
+
+
+def _top_dir(rel_path: str) -> str:
+    """Top-level directory of a repo-relative path ('.' for root files)."""
+    return rel_path.split("/", 1)[0] if "/" in rel_path else "."
+
+
+def extract_artifacts(
+    root: Path,
+    files: Sequence[scanner.ScannedFile],
+    *,
+    token_budget: int | None = None,
+    count_tokens: Callable[[str], int] | None = None,
+) -> ExtractionResult:
+    """Extract artifacts from text files within a token budget.
+
+    ``token_budget=None`` auto-scales: the whole corpus is kept up to a
+    1M-token ceiling. When the corpus exceeds the budget, artifacts are
+    allocated by importance instead of directory-walk order: parsed code
+    (functions/classes) beats whole-file excerpts, and the budget is spread
+    round-robin across top-level directories so one big folder can't starve
+    the rest. Drops are counted per directory — never silent.
+    """
+    if count_tokens is None:
+        from receipts.llm.token_estimate import estimate_tokens
+
+        count_tokens = estimate_tokens
+
+    # Phase 1 — collect every candidate artifact (no budget gate yet).
+    candidates: list[Artifact] = []
+    redaction_count = 0
+    for file in files:
+        if not file.is_text:
+            continue
+        text = scanner.read_text(root, file.rel_path)
+        if not text.strip():
+            continue
+
+        suffix = "." + file.rel_path.rsplit(".", 1)[-1] if "." in file.rel_path else ""
+        language = _EXTENSION_LANGUAGES.get(suffix)
+        raw: list[tuple[str, str, int, int, str]] = []
+        if language is not None:
+            try:
+                raw = _code_artifacts(file.rel_path, language, text)
+            except Exception:
+                raw = []
+        if not raw:
+            # Unsupported language, parse failure, or a script with no
+            # top-level defs — keep an excerpt rather than dropping the file.
+            raw = [_fallback_artifact(file.rel_path, text)]
+            language = language or "text"
+
+        for kind, name, start_line, end_line, code in raw:
+            clean_code, redactions = redact(code)
+            redaction_count += redactions
+            candidates.append(
+                Artifact(
+                    rel_path=file.rel_path,
+                    language=language,
+                    kind=kind,
+                    name=name,
+                    start_line=start_line,
+                    end_line=end_line,
+                    code=clean_code,
+                    token_estimate=count_tokens(clean_code),
+                )
+            )
+
+    corpus_tokens = sum(a.token_estimate for a in candidates)
+    auto = token_budget is None
+    budget = min(corpus_tokens, _AUTO_BUDGET_CEILING) if auto else token_budget
+
+    result = ExtractionResult(
+        budget=budget,
+        auto_budget=auto,
+        corpus_tokens=corpus_tokens,
+        redaction_count=redaction_count,
+    )
+
+    if corpus_tokens <= budget:
+        result.artifacts = candidates
+        result.used_tokens = corpus_tokens
+        return result
+
+    # Phase 2 — importance allocation. Tier 0: parsed code artifacts.
+    # Tier 1: whole-file excerpts (config, docs, unsupported languages).
+    # Within a tier, round-robin across top-level directories.
+    queues: dict[tuple[int, str], deque[Artifact]] = {}
+    dir_order: dict[int, list[str]] = {0: [], 1: []}
+    for a in candidates:
+        tier = 0 if a.kind != "file" else 1
+        key = (tier, _top_dir(a.rel_path))
+        if key not in queues:
+            queues[key] = deque()
+            dir_order[tier].append(_top_dir(a.rel_path))
+        queues[key].append(a)
+
+    for tier in (0, 1):
+        active = [queues[(tier, d)] for d in dir_order[tier]]
+        while any(active):
+            for q in active:
+                if not q:
+                    continue
+                a = q.popleft()
+                if result.used_tokens + a.token_estimate > budget:
+                    result.dropped_over_budget += 1
+                    result.dropped_tokens += a.token_estimate
+                    d = _top_dir(a.rel_path)
+                    result.dropped_by_dir[d] = result.dropped_by_dir.get(d, 0) + 1
+                    continue
+                result.used_tokens += a.token_estimate
+                result.artifacts.append(a)
+    return result