Add tree-sitter bash and JSON extractors (#866)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

graphify/extract.py

@@ -5579,6 +5579,295 @@ def _check_tree_sitter_version() -> None:
         )
 
 
+def extract_bash(path: Path) -> dict:
+    """Extract functions, source imports, and cross-function calls from a .sh file."""
+    try:
+        import tree_sitter_bash as tsbash
+        from tree_sitter import Language, Parser
+    except ImportError:
+        return {"nodes": [], "edges": [], "error": "tree-sitter-bash not installed"}
+
+    try:
+        language = Language(tsbash.language())
+        parser = Parser(language)
+        source = path.read_bytes()
+        tree = parser.parse(source)
+        root = tree.root_node
+    except Exception as e:
+        return {"nodes": [], "edges": [], "error": str(e)}
+
+    stem = _file_stem(path)
+    str_path = str(path)
+    nodes: list[dict] = []
+    edges: list[dict] = []
+    seen_ids: set[str] = set()
+    function_bodies: list[tuple[str, Any]] = []
+    defined_functions: set[str] = set()
+
+    def add_node(nid: str, label: str, line: int) -> None:
+        if nid and nid not in seen_ids:
+            seen_ids.add(nid)
+            nodes.append({"id": nid, "label": label, "file_type": "code",
+                          "source_file": str_path, "source_location": f"L{line}"})
+
+    def add_edge(src: str, tgt: str, relation: str, line: int,
+                 confidence: str = "EXTRACTED", weight: float = 1.0,
+                 context: str | None = None) -> None:
+        if not src or not tgt or src == tgt:
+            return
+        edge = {"source": src, "target": tgt, "relation": relation,
+                "confidence": confidence, "source_file": str_path,
+                "source_location": f"L{line}", "weight": weight}
+        if context:
+            edge["context"] = context
+        edges.append(edge)
+
+    file_nid = _make_id(str(path))
+    add_node(file_nid, path.name, 1)
+
+    _BASH_SKIP = frozenset({
+        "if", "then", "else", "elif", "fi", "for", "while", "until", "do",
+        "done", "case", "esac", "in", "return", "exit", "break", "continue",
+        "echo", "printf", "cd", "set", "local", "export", "readonly",
+        "declare", "unset", "shift", "read", "test", "[", "[[", ":", "true",
+        "false", "source", ".", "trap", "wait", "exec", "eval",
+    })
+
+    def _bash_func_name(node) -> str | None:
+        """Get the name from a function_definition node."""
+        # bash grammar: function_definition has a word child (the name)
+        for child in node.children:
+            if child.type == "word":
+                return _read_text(child, source)
+        return None
+
+    def walk_calls(body_node, func_nid: str, seen_calls: set) -> None:
+        if body_node is None:
+            return
+        for child in body_node.children:
+            if child.type == "command":
+                cmd_name_node = child.child_by_field_name("name")
+                if cmd_name_node is None and child.children:
+                    cmd_name_node = child.children[0]
+                if cmd_name_node:
+                    name = _read_text(cmd_name_node, source).strip()
+                    if name and name not in _BASH_SKIP and name in defined_functions:
+                        tgt = _make_id(stem, name)
+                        key = (func_nid, tgt)
+                        if tgt and key not in seen_calls:
+                            seen_calls.add(key)
+                            add_edge(func_nid, tgt, "calls",
+                                     child.start_point[0] + 1,
+                                     confidence="EXTRACTED", context="call")
+            walk_calls(child, func_nid, seen_calls)
+
+    def walk(node, parent_nid: str) -> None:
+        t = node.type
+        if t == "function_definition":
+            name = _bash_func_name(node)
+            if name:
+                fn_nid = _make_id(stem, name)
+                line = node.start_point[0] + 1
+                add_node(fn_nid, f"{name}()", line)
+                add_edge(parent_nid, fn_nid, "defines", line)
+                defined_functions.add(name)
+                # find the compound_statement body
+                body = None
+                for child in node.children:
+                    if child.type == "compound_statement":
+                        body = child
+                        break
+                function_bodies.append((fn_nid, body))
+            return  # don't recurse into function body during structural pass
+
+        if t == "command":
+            cmd_name_node = node.child_by_field_name("name")
+            if cmd_name_node is None and node.children:
+                cmd_name_node = node.children[0]
+            if cmd_name_node:
+                cmd = _read_text(cmd_name_node, source).strip()
+                if cmd in ("source", "."):
+                    # find the path argument (first word after command name)
+                    args = [c for c in node.children
+                            if c.type in ("word", "string", "concatenation")
+                            and c != cmd_name_node]
+                    if args:
+                        raw = _read_text(args[0], source).strip().strip("'\"")
+                        line = node.start_point[0] + 1
+                        if raw.startswith((".", "/")):
+                            resolved = (path.parent / raw).resolve()
+                            tgt_nid = _make_id(str(resolved))
+                            add_edge(file_nid, tgt_nid, "imports_from", line,
+                                     context="import")
+                        else:
+                            tgt_nid = _make_id(raw)
+                            if tgt_nid:
+                                add_edge(file_nid, tgt_nid, "imports", line,
+                                         context="import")
+            return
+
+        if t == "declaration_command":
+            # export/declare/readonly VAR=value at program level
+            if node.parent and node.parent.type == "program":
+                for child in node.children:
+                    if child.type == "variable_assignment":
+                        var_node = child.child_by_field_name("name")
+                        if var_node:
+                            var = _read_text(var_node, source).strip()
+                            if var:
+                                var_nid = _make_id(stem, var)
+                                line = child.start_point[0] + 1
+                                add_node(var_nid, var, line)
+                                add_edge(file_nid, var_nid, "defines", line)
+            return
+
+        for child in node.children:
+            walk(child, parent_nid)
+
+    walk(root, file_nid)
+
+    # Second pass: cross-function calls
+    for fn_nid, body in function_bodies:
+        walk_calls(body, fn_nid, set())
+
+    return {"nodes": nodes, "edges": edges}
+
+
+def extract_json(path: Path) -> dict:
+    """Extract top-level keys, nested structure, and dependency edges from a .json file."""
+    _JSON_MAX_BYTES = 1_048_576  # 1 MiB — skip large fixture dumps / GeoJSON blobs
+
+    try:
+        import tree_sitter_json as tsjson
+        from tree_sitter import Language, Parser
+    except ImportError:
+        return {"nodes": [], "edges": [], "error": "tree-sitter-json not installed"}
+
+    try:
+        if path.stat().st_size > _JSON_MAX_BYTES:
+            return {"nodes": [], "edges": [], "error": "json file too large to index"}
+        language = Language(tsjson.language())
+        parser = Parser(language)
+        source = path.read_bytes()
+        tree = parser.parse(source)
+        root = tree.root_node
+    except Exception as e:
+        return {"nodes": [], "edges": [], "error": str(e)}
+
+    stem = _file_stem(path)
+    str_path = str(path)
+    nodes: list[dict] = []
+    edges: list[dict] = []
+    seen_ids: set[str] = set()
+
+    # Keys whose string values become imports (package.json dep blocks)
+    _DEP_KEYS = frozenset({
+        "dependencies", "devDependencies", "peerDependencies",
+        "optionalDependencies", "bundleDependencies", "bundledDependencies",
+    })
+
+    def add_node(nid: str, label: str, line: int) -> None:
+        if nid and nid not in seen_ids:
+            seen_ids.add(nid)
+            nodes.append({"id": nid, "label": label, "file_type": "code",
+                          "source_file": str_path, "source_location": f"L{line}"})
+
+    def add_edge(src: str, tgt: str, relation: str, line: int,
+                 context: str | None = None) -> None:
+        if not src or not tgt or src == tgt:
+            return
+        edge = {"source": src, "target": tgt, "relation": relation,
+                "confidence": "EXTRACTED", "source_file": str_path,
+                "source_location": f"L{line}", "weight": 1.0}
+        if context:
+            edge["context"] = context
+        edges.append(edge)
+
+    file_nid = _make_id(str(path))
+    add_node(file_nid, path.name, 1)
+
+    def _key_text(pair_node) -> str | None:
+        """Extract the string content of a pair's key."""
+        key_node = pair_node.child_by_field_name("key")
+        if key_node is None:
+            return None
+        if key_node.type == "string":
+            content = key_node.child_by_field_name("string_content")
+            if content:
+                return _read_text(content, source)
+            # fallback: strip surrounding quotes
+            raw = _read_text(key_node, source)
+            return raw.strip('"\'')
+        return _read_text(key_node, source)
+
+    def _val_node(pair_node):
+        return pair_node.child_by_field_name("value")
+
+    def walk_object(obj_node, parent_nid: str, parent_key: str | None,
+                    depth: int, pair_count: list) -> None:
+        if depth > 6 or pair_count[0] > 500:
+            return
+        for child in obj_node.children:
+            if child.type != "pair":
+                continue
+            pair_count[0] += 1
+            key = _key_text(child)
+            if not key:
+                continue
+            key_nid = _make_id(stem, *(([parent_key] if parent_key else []) + [key]))
+            if not key_nid:
+                continue
+            line = child.start_point[0] + 1
+            add_node(key_nid, key, line)
+            add_edge(parent_nid, key_nid, "contains", line)
+
+            val = _val_node(child)
+            if val is None:
+                continue
+
+            if val.type == "object":
+                walk_object(val, key_nid, key, depth + 1, pair_count)
+
+            elif val.type == "array":
+                # For "extends" arrays (tsconfig, eslint): each string element
+                for item in val.children:
+                    if item.type == "string":
+                        content = item.child_by_field_name("string_content")
+                        ref = _read_text(content, source) if content else _read_text(item, source).strip('"\'')
+                        if ref:
+                            ref_nid = _make_id(ref)
+                            if ref_nid:
+                                add_edge(key_nid, ref_nid, "extends", line, context="import")
+
+            elif val.type == "string":
+                content = val.child_by_field_name("string_content")
+                val_text = _read_text(content, source) if content else _read_text(val, source).strip('"\'')
+
+                if key == "extends" and val_text:
+                    ref_nid = _make_id(val_text)
+                    if ref_nid:
+                        add_edge(file_nid, ref_nid, "extends", line, context="import")
+
+                elif key == "$ref" and val_text:
+                    ref_nid = _make_id(val_text)
+                    if ref_nid:
+                        add_edge(parent_nid, ref_nid, "references", line)
+
+                elif parent_key in _DEP_KEYS and val_text:
+                    dep_nid = _make_id(key)
+                    if dep_nid:
+                        add_edge(key_nid, dep_nid, "imports", line, context="import")
+
+    # Entry: find root document → object
+    doc = root
+    if doc.type == "document" and doc.child_count > 0:
+        doc = doc.children[0]
+    if doc.type == "object":
+        walk_object(doc, file_nid, None, 0, [0])
+
+    return {"nodes": nodes, "edges": edges}
+
+
 _DISPATCH: dict[str, Any] = {
     ".py": extract_python,
     ".js": extract_js,
@@ -5643,6 +5932,9 @@ _DISPATCH: dict[str, Any] = {
     ".dfm": extract_delphi_form,
     ".lfm": extract_lazarus_form,
     ".lpk": extract_lazarus_package,
+    ".sh": extract_bash,
+    ".bash": extract_bash,
+    ".json": extract_json,
 }
 
 

pyproject.toml

@@ -38,6 +38,8 @@ dependencies = [
     "tree-sitter-julia",
     "tree-sitter-verilog",
     "tree-sitter-fortran",
+    "tree-sitter-bash",
+    "tree-sitter-json",
 ]
 
 [project.urls]

tests/fixtures/sample.json

@@ -0,0 +1,16 @@
+{
+  "name": "my-app",
+  "version": "1.0.0",
+  "scripts": {
+    "build": "tsc",
+    "test": "jest",
+    "start": "node dist/index.js"
+  },
+  "dependencies": {
+    "react": "^18.0.0",
+    "axios": "^1.6.0"
+  },
+  "devDependencies": {
+    "typescript": "^5.0.0"
+  }
+}

tests/fixtures/sample.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+source ./helpers.sh
+
+export APP_ENV="production"
+
+build() {
+    echo "Building..."
+    local out_dir="dist"
+    mkdir -p "$out_dir"
+}
+
+test_suite() {
+    echo "Running tests..."
+    build
+}
+
+deploy() {
+    build
+    test_suite
+    echo "Deploying to $APP_ENV"
+}

tests/fixtures/sample_tsconfig.json

@@ -0,0 +1,10 @@
+{
+  "extends": "@tsconfig/strictest/tsconfig.json",
+  "compilerOptions": {
+    "target": "ES2022",
+    "module": "NodeNext",
+    "outDir": "dist",
+    "strict": true
+  },
+  "include": ["src"]
+}

tests/test_extract.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from graphify.extract import extract_python, extract, collect_files, _make_id
+from graphify.extract import extract_python, extract, collect_files, _make_id, extract_bash, extract_json, _DISPATCH
 
 FIXTURES = Path(__file__).parent / "fixtures"
 
@@ -425,3 +425,170 @@ def test_extract_parallel_returns_false_on_broken_pool(tmp_path, monkeypatch, ca
     out = capsys.readouterr().out
     assert "BrokenProcessPool" in out, "user-facing warning must mention the failure"
     assert "__main__" in out, "warning must hint at the Windows __main__ guard idiom"
+
+
+# ---------------------------------------------------------------------------
+# Bash extractor tests (#866)
+# ---------------------------------------------------------------------------
+
+def test_dispatch_includes_sh_and_json():
+    assert ".sh" in _DISPATCH
+    assert ".bash" in _DISPATCH
+    assert ".json" in _DISPATCH
+
+
+def test_extract_bash_finds_functions():
+    result = extract_bash(FIXTURES / "sample.sh")
+    assert "error" not in result
+    labels = {n["label"] for n in result["nodes"]}
+    assert "build()" in labels
+    assert "test_suite()" in labels
+    assert "deploy()" in labels
+
+
+def test_extract_bash_emits_defines_edges():
+    result = extract_bash(FIXTURES / "sample.sh")
+    relations = {e["relation"] for e in result["edges"]}
+    assert "defines" in relations
+
+
+def test_extract_bash_emits_calls_edges():
+    result = extract_bash(FIXTURES / "sample.sh")
+    calls = [(e["source"], e["target"]) for e in result["edges"] if e["relation"] == "calls"]
+    # deploy() calls build() and test_suite(); test_suite() calls build()
+    assert any("deploy" in s and "build" in t for s, t in calls)
+    assert any("deploy" in s and "test_suite" in t for s, t in calls)
+    assert any("test_suite" in s and "build" in t for s, t in calls)
+
+
+def test_extract_bash_calls_have_extracted_confidence():
+    result = extract_bash(FIXTURES / "sample.sh")
+    for e in result["edges"]:
+        if e["relation"] == "calls":
+            assert e["confidence"] == "EXTRACTED"
+            assert e.get("context") == "call"
+
+
+def test_extract_bash_emits_source_imports_from(tmp_path):
+    helpers = tmp_path / "helpers.sh"
+    helpers.write_text("# helper\n")
+    script = tmp_path / "deploy.sh"
+    script.write_text(f"#!/bin/bash\nsource ./helpers.sh\nfoo() {{ echo hi; }}\n")
+    result = extract_bash(script)
+    import_edges = [e for e in result["edges"] if e["relation"] == "imports_from"]
+    assert len(import_edges) >= 1
+    assert import_edges[0].get("context") == "import"
+
+
+def test_extract_bash_no_self_loops():
+    result = extract_bash(FIXTURES / "sample.sh")
+    for e in result["edges"]:
+        assert e["source"] != e["target"], f"Self-loop: {e}"
+
+
+def test_extract_bash_no_dangling_edges():
+    result = extract_bash(FIXTURES / "sample.sh")
+    node_ids = {n["id"] for n in result["nodes"]}
+    for e in result["edges"]:
+        assert e["source"] in node_ids, f"Dangling source: {e['source']}"
+        # targets may reference external files (imports_from) — only check non-import edges
+        if e["relation"] not in ("imports_from", "imports"):
+            assert e["target"] in node_ids, f"Dangling target: {e['target']}"
+
+
+def test_extract_bash_skip_builtins_in_calls():
+    result = extract_bash(FIXTURES / "sample.sh")
+    builtins = {"echo", "cd", "set", "export", "local", "mkdir", "if", "then"}
+    call_targets = {e["target"] for e in result["edges"] if e["relation"] == "calls"}
+    for b in builtins:
+        assert not any(b in t for t in call_targets), f"Builtin '{b}' appeared as calls target"
+
+
+def test_extract_bash_missing_grammar_returns_error():
+    """extract_bash returns error dict when tree-sitter-bash not installed (mocked)."""
+    import unittest.mock as mock
+    import builtins
+    real_import = builtins.__import__
+
+    def patched(name, *args, **kwargs):
+        if name == "tree_sitter_bash":
+            raise ImportError("mocked")
+        return real_import(name, *args, **kwargs)
+
+    with mock.patch("builtins.__import__", side_effect=patched):
+        result = extract_bash(FIXTURES / "sample.sh")
+    assert "error" in result
+    assert result["nodes"] == []
+
+
+# ---------------------------------------------------------------------------
+# JSON extractor tests (#866)
+# ---------------------------------------------------------------------------
+
+def test_extract_json_top_level_keys():
+    result = extract_json(FIXTURES / "sample.json")
+    assert "error" not in result
+    labels = {n["label"] for n in result["nodes"]}
+    assert "name" in labels
+    assert "version" in labels
+    assert "scripts" in labels
+    assert "dependencies" in labels
+
+
+def test_extract_json_nested_contains():
+    result = extract_json(FIXTURES / "sample.json")
+    contains = [(e["source"], e["target"]) for e in result["edges"] if e["relation"] == "contains"]
+    assert any("scripts" in s and "build" in t for s, t in contains)
+    assert any("scripts" in s and "test" in t for s, t in contains)
+    assert any("dependencies" in s and "react" in t for s, t in contains)
+
+
+def test_extract_json_dependencies_become_imports():
+    result = extract_json(FIXTURES / "sample.json")
+    import_edges = [e for e in result["edges"] if e["relation"] == "imports"]
+    targets = {e["target"] for e in import_edges}
+    assert any("react" in t for t in targets)
+    assert any("axios" in t for t in targets)
+    assert any("typescript" in t for t in targets)
+
+
+def test_extract_json_extends_resolved():
+    result = extract_json(FIXTURES / "sample_tsconfig.json")
+    extends_edges = [e for e in result["edges"] if e["relation"] == "extends"]
+    assert len(extends_edges) >= 1
+    assert extends_edges[0].get("context") == "import"
+
+
+def test_extract_json_large_file_skipped(tmp_path):
+    big = tmp_path / "big.json"
+    # Write a JSON file just over 1 MiB
+    big.write_bytes(b'{"x": "' + b"a" * (1_048_576) + b'"}')
+    result = extract_json(big)
+    assert "error" in result
+    assert result["nodes"] == []
+
+
+def test_extract_json_handles_invalid_json(tmp_path):
+    bad = tmp_path / "broken.json"
+    bad.write_text("{this is not: valid json!!!")
+    result = extract_json(bad)
+    # Should not crash — returns empty or error result
+    assert isinstance(result, dict)
+    assert "nodes" in result
+
+
+def test_extract_json_no_self_loops():
+    result = extract_json(FIXTURES / "sample.json")
+    for e in result["edges"]:
+        assert e["source"] != e["target"], f"Self-loop: {e}"
+
+
+def test_extract_bash_via_dispatch():
+    from graphify.extract import _get_extractor
+    assert _get_extractor(Path("foo.sh")) is extract_bash
+    assert _get_extractor(Path("foo.bash")) is extract_bash
+
+
+def test_extract_json_via_dispatch():
+    from graphify.extract import _get_extractor
+    assert _get_extractor(Path("foo.json")) is extract_json