analyzer/js_treesitter.py

"""Extract structure from JS / TS files using tree-sitter.

We use `tree-sitter-languages`, which ships prebuilt grammars so the Space needs
no compiler. If tree-sitter is unavailable we fall back to a small regex pass so
the app still works (degraded, but never crashes).
"""
from __future__ import annotations

import re

from schema import Symbol

try:  # prebuilt grammars; no compiler needed on the Space
    from tree_sitter_languages import get_parser
    _HAVE_TS = True
except Exception:  # pragma: no cover - exercised only when dep missing
    _HAVE_TS = False

_LANG_NAME = {"javascript": "javascript", "typescript": "typescript"}


def parse(text: str, lang: str) -> tuple[list[Symbol], list[str]]:
    """Return (symbols, import_targets). Best-effort; never raises."""
    if _HAVE_TS:
        try:
            return _parse_ts(text, lang)
        except Exception:
            pass
    return _parse_regex(text)


# --- tree-sitter path -------------------------------------------------------
def _parse_ts(text: str, lang: str) -> tuple[list[Symbol], list[str]]:
    parser = get_parser(_LANG_NAME.get(lang, "javascript"))
    blob = text.encode("utf-8", "ignore")
    root = parser.parse(blob).root_node
    symbols: list[Symbol] = []
    imports: list[str] = []

    def txt(node) -> str:
        return blob[node.start_byte:node.end_byte].decode("utf-8", "ignore")

    def name_of(node) -> str | None:
        n = node.child_by_field_name("name")
        return txt(n) if n is not None else None

    def visit(node, depth=0):
        t = node.type
        if t in ("function_declaration", "generator_function_declaration"):
            nm = name_of(node)
            if nm:
                symbols.append(Symbol(kind="function", name=nm, lineno=node.start_point[0] + 1))
        elif t in ("class_declaration", "abstract_class_declaration"):
            nm = name_of(node)
            if nm:
                symbols.append(Symbol(kind="class", name=nm, lineno=node.start_point[0] + 1))
        elif t == "lexical_declaration":
            # const Foo = () => {...}  /  const Foo = function...
            for child in node.named_children:
                if child.type == "variable_declarator":
                    val = child.child_by_field_name("value")
                    nm = name_of(child)
                    if nm and val is not None and val.type in (
                            "arrow_function", "function", "function_expression"):
                        symbols.append(Symbol(kind="function", name=nm,
                                              lineno=node.start_point[0] + 1))
        elif t == "import_statement":
            src = node.child_by_field_name("source")
            if src is not None:
                imports.append(txt(src).strip("'\"`"))
        for child in node.named_children:
            visit(child, depth + 1)

    visit(root)
    # require('x') and dynamic import('x') aren't import_statements — sweep them too.
    imports.extend(_REQUIRE_RE.findall(text))
    seen: set[str] = set()
    uniq = [i for i in imports if i and not (i in seen or seen.add(i))]
    return symbols, uniq


# --- regex fallback ---------------------------------------------------------
_FUNC_RE = re.compile(r"\bfunction\s+([A-Za-z_$][\w$]*)")
_ARROW_RE = re.compile(r"\b(?:const|let|var)\s+([A-Za-z_$][\w$]*)\s*=\s*(?:async\s*)?\([^)]*\)\s*=>")
_CLASS_RE = re.compile(r"\bclass\s+([A-Za-z_$][\w$]*)")
_IMPORT_RE = re.compile(r"""\bimport\b[^'"]*['"]([^'"]+)['"]""")
_REQUIRE_RE = re.compile(r"""\brequire\(\s*['"]([^'"]+)['"]\s*\)""")


def _parse_regex(text: str) -> tuple[list[Symbol], list[str]]:
    symbols: list[Symbol] = []
    for m in _FUNC_RE.finditer(text):
        symbols.append(Symbol(kind="function", name=m.group(1)))
    for m in _ARROW_RE.finditer(text):
        symbols.append(Symbol(kind="function", name=m.group(1)))
    for m in _CLASS_RE.finditer(text):
        symbols.append(Symbol(kind="class", name=m.group(1)))
    imports = _IMPORT_RE.findall(text) + _REQUIRE_RE.findall(text)
    seen: set[str] = set()
    uniq = [i for i in imports if not (i in seen or seen.add(i))]
    return symbols, uniq