From a752082442e390c1d2fa36652c10cf26cce77c24 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Hendrik=20Eisenbl=C3=A4tter?= Date: Mon, 6 Apr 2026 23:57:55 +0200 Subject: [PATCH] support zig and powershell --- graphify/detect.py | 2 +- graphify/extract.py | 395 ++++++++++++++++++++++++++++++++++++++ pyproject.toml | 2 + tests/fixtures/sample.ps1 | 37 ++++ tests/fixtures/sample.zig | 36 ++++ tests/test_extract.py | 3 +- tests/test_languages.py | 79 +++++++- 7 files changed, 551 insertions(+), 3 deletions(-) create mode 100644 tests/fixtures/sample.ps1 create mode 100644 tests/fixtures/sample.zig diff --git a/graphify/detect.py b/graphify/detect.py index 3c7406886..607b85818 100644 --- a/graphify/detect.py +++ b/graphify/detect.py @@ -16,7 +16,7 @@ class FileType(str, Enum): _MANIFEST_PATH = "graphify-out/manifest.json" -CODE_EXTENSIONS = {'.py', '.ts', '.js', '.tsx', '.go', '.rs', '.java', '.cpp', '.cc', '.cxx', '.c', '.h', '.hpp', '.rb', '.swift', '.kt', '.kts', '.cs', '.scala', '.php'} +CODE_EXTENSIONS = {'.py', '.ts', '.js', '.tsx', '.go', '.rs', '.java', '.cpp', '.cc', '.cxx', '.c', '.h', '.hpp', '.rb', '.swift', '.kt', '.kts', '.cs', '.scala', '.php', '.zig', '.ps1'} DOC_EXTENSIONS = {'.md', '.txt', '.rst'} PAPER_EXTENSIONS = {'.pdf'} IMAGE_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.gif', '.webp', '.svg'} diff --git a/graphify/extract.py b/graphify/extract.py index 70c468e3e..6adc3bcb5 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -1341,6 +1341,398 @@ def walk_calls(node, caller_nid: str) -> None: return {"nodes": nodes, "edges": clean_edges} +# ── Zig ────────────────────────────────────────────────────────────────────── + +def extract_zig(path: Path) -> dict: + """Extract functions, structs, enums, unions, and imports from a .zig file.""" + try: + import tree_sitter_zig as tszig + from tree_sitter import Language, Parser + except ImportError: + return {"nodes": [], "edges": [], "error": "tree-sitter-zig not installed"} + + try: + language = Language(tszig.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 = path.stem + str_path = str(path) + nodes: list[dict] = [] + edges: list[dict] = [] + seen_ids: set[str] = set() + function_bodies: list[tuple[str, object]] = [] + + def add_node(nid: str, label: str, line: int) -> None: + if 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) -> None: + edges.append({ + "source": src, + "target": tgt, + "relation": relation, + "confidence": confidence, + "source_file": str_path, + "source_location": f"L{line}", + "weight": weight, + }) + + file_nid = _make_id(stem) + add_node(file_nid, path.name, 1) + + def _extract_import(node) -> None: + """Handle @import / @cImport builtin calls inside variable_declaration.""" + for child in node.children: + if child.type == "builtin_function": + bi = None + args = None + for c in child.children: + if c.type == "builtin_identifier": + bi = _read_text(c, source) + elif c.type == "arguments": + args = c + if bi in ("@import", "@cImport") and args: + for arg in args.children: + if arg.type in ("string_literal", "string"): + raw = _read_text(arg, source).strip('"') + module_name = raw.split("/")[-1].split(".")[0] + if module_name: + tgt_nid = _make_id(module_name) + add_edge(file_nid, tgt_nid, "imports_from", + node.start_point[0] + 1) + return + elif child.type == "field_expression": + # const mem = @import("std").mem — recurse into field_expression + _extract_import(child) + return + + def walk(node, parent_struct_nid: str | None = None) -> None: + t = node.type + + if t == "function_declaration": + name_node = node.child_by_field_name("name") + if name_node: + func_name = _read_text(name_node, source) + line = node.start_point[0] + 1 + if parent_struct_nid: + func_nid = _make_id(parent_struct_nid, func_name) + add_node(func_nid, f".{func_name}()", line) + add_edge(parent_struct_nid, func_nid, "method", line) + else: + func_nid = _make_id(stem, func_name) + add_node(func_nid, f"{func_name}()", line) + add_edge(file_nid, func_nid, "contains", line) + body = node.child_by_field_name("body") + if body: + function_bodies.append((func_nid, body)) + return + + if t == "variable_declaration": + # Zig overloads variable_declaration for structs, enums, unions, imports + name_node = None + value_node = None + for child in node.children: + if child.type == "identifier": + name_node = child + elif child.type in ("struct_declaration", "enum_declaration", + "union_declaration", "builtin_function", + "field_expression"): + value_node = child + + if value_node and value_node.type == "struct_declaration": + if name_node: + struct_name = _read_text(name_node, source) + line = node.start_point[0] + 1 + struct_nid = _make_id(stem, struct_name) + add_node(struct_nid, struct_name, line) + add_edge(file_nid, struct_nid, "contains", line) + for child in value_node.children: + walk(child, parent_struct_nid=struct_nid) + return + + if value_node and value_node.type in ("enum_declaration", "union_declaration"): + if name_node: + type_name = _read_text(name_node, source) + line = node.start_point[0] + 1 + type_nid = _make_id(stem, type_name) + add_node(type_nid, type_name, line) + add_edge(file_nid, type_nid, "contains", line) + return + + if value_node and value_node.type in ("builtin_function", "field_expression"): + _extract_import(node) + return + + return + + for child in node.children: + walk(child, parent_struct_nid) + + walk(root) + + # Call-graph pass + label_to_nid: dict[str, str] = {} + for n in nodes: + raw = n["label"] + normalised = raw.strip("()").lstrip(".") + label_to_nid[normalised.lower()] = n["id"] + + seen_call_pairs: set[tuple[str, str]] = set() + + def walk_calls(node, caller_nid: str) -> None: + if node.type == "function_declaration": + return + if node.type == "call_expression": + func_node = node.child_by_field_name("function") + callee_name: str | None = None + if func_node: + if func_node.type == "identifier": + callee_name = _read_text(func_node, source) + elif func_node.type == "field_expression": + member = func_node.child_by_field_name("member") + if member: + callee_name = _read_text(member, source) + if callee_name: + tgt_nid = label_to_nid.get(callee_name.lower()) + if tgt_nid and tgt_nid != caller_nid: + pair = (caller_nid, tgt_nid) + if pair not in seen_call_pairs: + seen_call_pairs.add(pair) + add_edge(caller_nid, tgt_nid, "calls", + node.start_point[0] + 1, + confidence="INFERRED", weight=0.8) + for child in node.children: + walk_calls(child, caller_nid) + + for caller_nid, body_node in function_bodies: + walk_calls(body_node, caller_nid) + + valid_ids = seen_ids + clean_edges = [] + for edge in edges: + src, tgt = edge["source"], edge["target"] + if src in valid_ids and (tgt in valid_ids or edge["relation"] in ("imports", "imports_from")): + clean_edges.append(edge) + + return {"nodes": nodes, "edges": clean_edges} + + +# ── PowerShell ─────────────────────────────────────────────────────────────── + +def extract_powershell(path: Path) -> dict: + """Extract functions, classes, methods, using statements, and cmdlet calls from a .ps1 file.""" + try: + import tree_sitter_powershell as tsps + from tree_sitter import Language, Parser + except ImportError: + return {"nodes": [], "edges": [], "error": "tree-sitter-powershell not installed"} + + try: + language = Language(tsps.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 = path.stem + str_path = str(path) + nodes: list[dict] = [] + edges: list[dict] = [] + seen_ids: set[str] = set() + function_bodies: list[tuple[str, object]] = [] + + def add_node(nid: str, label: str, line: int) -> None: + if 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) -> None: + edges.append({ + "source": src, + "target": tgt, + "relation": relation, + "confidence": confidence, + "source_file": str_path, + "source_location": f"L{line}", + "weight": weight, + }) + + file_nid = _make_id(stem) + add_node(file_nid, path.name, 1) + + _PS_SKIP_COMMANDS = frozenset({ + "using", "return", "if", "else", "elseif", "foreach", "for", + "while", "do", "switch", "try", "catch", "finally", "throw", + "break", "continue", "exit", "param", "begin", "process", "end", + }) + + def _find_script_block_body(node): + """Find script_block_body inside a function/method's script_block.""" + for child in node.children: + if child.type == "script_block": + for sc in child.children: + if sc.type == "script_block_body": + return sc + return child + return None + + def walk(node, parent_class_nid: str | None = None) -> None: + t = node.type + + if t == "function_statement": + name_node = None + for child in node.children: + if child.type == "function_name": + name_node = child + break + if name_node: + func_name = _read_text(name_node, source) + line = node.start_point[0] + 1 + func_nid = _make_id(stem, func_name) + add_node(func_nid, f"{func_name}()", line) + add_edge(file_nid, func_nid, "contains", line) + body = _find_script_block_body(node) + if body: + function_bodies.append((func_nid, body)) + return + + if t == "class_statement": + name_node = None + for child in node.children: + if child.type == "simple_name": + name_node = child + break + if name_node: + class_name = _read_text(name_node, source) + line = node.start_point[0] + 1 + class_nid = _make_id(stem, class_name) + add_node(class_nid, class_name, line) + add_edge(file_nid, class_nid, "contains", line) + for child in node.children: + walk(child, parent_class_nid=class_nid) + return + + if t == "class_method_definition": + name_node = None + for child in node.children: + if child.type == "simple_name": + name_node = child + break + if name_node: + method_name = _read_text(name_node, source) + line = node.start_point[0] + 1 + if parent_class_nid: + method_nid = _make_id(parent_class_nid, method_name) + add_node(method_nid, f".{method_name}()", line) + add_edge(parent_class_nid, method_nid, "method", line) + else: + method_nid = _make_id(stem, method_name) + add_node(method_nid, f"{method_name}()", line) + add_edge(file_nid, method_nid, "contains", line) + body = _find_script_block_body(node) + if body: + function_bodies.append((method_nid, body)) + return + + if t == "command": + # Distinguish using statements from regular commands + cmd_name_node = None + for child in node.children: + if child.type == "command_name": + cmd_name_node = child + break + if cmd_name_node: + cmd_text = _read_text(cmd_name_node, source).lower() + if cmd_text == "using": + # Parse: using namespace X.Y or using module X + tokens = [] + for child in node.children: + if child.type == "command_elements": + for el in child.children: + if el.type == "generic_token": + tokens.append(_read_text(el, source)) + # Skip keyword (namespace/module), take last token + module_tokens = [t for t in tokens + if t.lower() not in ("namespace", "module", "assembly")] + if module_tokens: + module_name = module_tokens[-1].split(".")[-1] + tgt_nid = _make_id(module_name) + add_edge(file_nid, tgt_nid, "imports_from", + node.start_point[0] + 1) + return + + for child in node.children: + walk(child, parent_class_nid) + + walk(root) + + # Call-graph pass + label_to_nid: dict[str, str] = {} + for n in nodes: + raw = n["label"] + normalised = raw.strip("()").lstrip(".") + label_to_nid[normalised.lower()] = n["id"] + + seen_call_pairs: set[tuple[str, str]] = set() + + def walk_calls(node, caller_nid: str) -> None: + if node.type in ("function_statement", "class_statement"): + return + if node.type == "command": + cmd_name_node = None + for child in node.children: + if child.type == "command_name": + cmd_name_node = child + break + if cmd_name_node: + cmd_text = _read_text(cmd_name_node, source) + if cmd_text.lower() not in _PS_SKIP_COMMANDS: + tgt_nid = label_to_nid.get(cmd_text.lower()) + if tgt_nid and tgt_nid != caller_nid: + pair = (caller_nid, tgt_nid) + if pair not in seen_call_pairs: + seen_call_pairs.add(pair) + add_edge(caller_nid, tgt_nid, "calls", + node.start_point[0] + 1, + confidence="INFERRED", weight=0.8) + for child in node.children: + walk_calls(child, caller_nid) + + for caller_nid, body_node in function_bodies: + walk_calls(body_node, caller_nid) + + valid_ids = seen_ids + clean_edges = [] + for edge in edges: + src, tgt = edge["source"], edge["target"] + if src in valid_ids and (tgt in valid_ids or edge["relation"] in ("imports", "imports_from")): + clean_edges.append(edge) + + return {"nodes": nodes, "edges": clean_edges} + + # ── Cross-file import resolution ────────────────────────────────────────────── def _resolve_cross_file_imports( @@ -1523,6 +1915,8 @@ def extract(paths: list[Path]) -> dict: ".kts": extract_kotlin, ".scala": extract_scala, ".php": extract_php, + ".zig": extract_zig, + ".ps1": extract_powershell, } for path in paths: @@ -1565,6 +1959,7 @@ def collect_files(target: Path) -> list[Path]: "*.py", "*.js", "*.ts", "*.tsx", "*.go", "*.rs", "*.java", "*.c", "*.h", "*.cpp", "*.cc", "*.cxx", "*.hpp", "*.rb", "*.cs", "*.kt", "*.kts", "*.scala", "*.php", + "*.zig", "*.ps1", ) results: list[Path] = [] for pattern in _EXTENSIONS: diff --git a/pyproject.toml b/pyproject.toml index 5a4a1e3cd..4f174cb54 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,8 @@ dependencies = [ "tree-sitter-kotlin", "tree-sitter-scala", "tree-sitter-php", + "tree-sitter-zig", + "tree-sitter-powershell", ] [project.urls] diff --git a/tests/fixtures/sample.ps1 b/tests/fixtures/sample.ps1 new file mode 100644 index 000000000..ec8a7ad64 --- /dev/null +++ b/tests/fixtures/sample.ps1 @@ -0,0 +1,37 @@ +using namespace System.IO +using module MyModule + +function Get-Data { + param( + [string]$Name, + [int]$Count = 10 + ) + + $result = Process-Items -Name $Name -Count $Count + return $result +} + +function Process-Items { + param( + [string]$Name, + [int]$Count + ) + + return @($Name) * $Count +} + +class DataProcessor { + [string]$Name + + DataProcessor([string]$name) { + $this.Name = $name + } + + [void] Process() { + $data = Get-Data -Name $this.Name + } + + [string] ToString() { + return $this.Name + } +} diff --git a/tests/fixtures/sample.zig b/tests/fixtures/sample.zig new file mode 100644 index 000000000..2df2955dc --- /dev/null +++ b/tests/fixtures/sample.zig @@ -0,0 +1,36 @@ +const std = @import("std"); +const mem = @import("std").mem; + +const Config = struct { + name: []const u8, + value: u32, + + pub fn init(name: []const u8, value: u32) Config { + return Config{ .name = name, .value = value }; + } + + pub fn getName(self: Config) []const u8 { + return self.name; + } +}; + +const Status = enum { + active, + inactive, +}; + +fn processData(cfg: Config) void { + const name = cfg.getName(); + _ = name; +} + +fn createConfig() Config { + const cfg = Config.init("test", 42); + processData(cfg); + return cfg; +} + +pub fn main() void { + const cfg = createConfig(); + _ = cfg; +} diff --git a/tests/test_extract.py b/tests/test_extract.py index 2ae63eaed..744eec491 100644 --- a/tests/test_extract.py +++ b/tests/test_extract.py @@ -60,7 +60,8 @@ def test_collect_files_from_dir(): files = collect_files(FIXTURES) supported = {".py", ".js", ".ts", ".tsx", ".go", ".rs", ".java", ".c", ".cpp", ".cc", ".cxx", ".rb", - ".cs", ".kt", ".kts", ".scala", ".php", ".h", ".hpp"} + ".cs", ".kt", ".kts", ".scala", ".php", ".h", ".hpp", + ".zig", ".ps1"} assert all(f.suffix in supported for f in files) assert len(files) > 0 diff --git a/tests/test_languages.py b/tests/test_languages.py index 3bd2d56ac..d0192d07c 100644 --- a/tests/test_languages.py +++ b/tests/test_languages.py @@ -1,10 +1,11 @@ -"""Tests for the 8 new language extractors: Java, C, C++, Ruby, C#, Kotlin, Scala, PHP.""" +"""Tests for language extractors: Java, C, C++, Ruby, C#, Kotlin, Scala, PHP, Zig, PowerShell.""" from __future__ import annotations from pathlib import Path import pytest from graphify.extract import ( extract_java, extract_c, extract_cpp, extract_ruby, extract_csharp, extract_kotlin, extract_scala, extract_php, + extract_zig, extract_powershell, ) FIXTURES = Path(__file__).parent / "fixtures" @@ -217,3 +218,79 @@ def test_php_finds_function(): def test_php_finds_imports(): r = extract_php(FIXTURES / "sample.php") assert "imports" in _relations(r) + + +# ── Zig ────────────────────────────────────────────────────────────────────── + +def test_zig_no_error(): + r = extract_zig(FIXTURES / "sample.zig") + assert "error" not in r + +def test_zig_finds_struct(): + r = extract_zig(FIXTURES / "sample.zig") + assert any("Config" in l for l in _labels(r)) + +def test_zig_finds_enum(): + r = extract_zig(FIXTURES / "sample.zig") + assert any("Status" in l for l in _labels(r)) + +def test_zig_finds_functions(): + r = extract_zig(FIXTURES / "sample.zig") + labels = _labels(r) + assert any("processData" in l for l in labels) + assert any("createConfig" in l for l in labels) + assert any("main" in l for l in labels) + +def test_zig_finds_struct_methods(): + r = extract_zig(FIXTURES / "sample.zig") + labels = _labels(r) + assert any(".init()" in l for l in labels) + assert any(".getName()" in l for l in labels) + +def test_zig_finds_imports(): + r = extract_zig(FIXTURES / "sample.zig") + assert "imports_from" in _relations(r) + +def test_zig_no_dangling_edges(): + r = extract_zig(FIXTURES / "sample.zig") + node_ids = {n["id"] for n in r["nodes"]} + for e in r["edges"]: + assert e["source"] in node_ids + + +# ── PowerShell ─────────────────────────────────────────────────────────────── + +def test_powershell_no_error(): + r = extract_powershell(FIXTURES / "sample.ps1") + assert "error" not in r + +def test_powershell_finds_class(): + r = extract_powershell(FIXTURES / "sample.ps1") + assert any("DataProcessor" in l for l in _labels(r)) + +def test_powershell_finds_functions(): + r = extract_powershell(FIXTURES / "sample.ps1") + labels = _labels(r) + assert any("Get-Data" in l for l in labels) + assert any("Process-Items" in l for l in labels) + +def test_powershell_finds_methods(): + r = extract_powershell(FIXTURES / "sample.ps1") + labels = _labels(r) + assert any(".Process()" in l for l in labels) + assert any(".ToString()" in l for l in labels) + +def test_powershell_finds_usings(): + r = extract_powershell(FIXTURES / "sample.ps1") + assert "imports_from" in _relations(r) + +def test_powershell_emits_calls(): + r = extract_powershell(FIXTURES / "sample.ps1") + calls = _calls(r) + assert any("Get-Data" in src for src, tgt in calls) + +def test_powershell_no_dangling_edges(): + r = extract_powershell(FIXTURES / "sample.ps1") + node_ids = {n["id"] for n in r["nodes"]} + for e in r["edges"]: + assert e["source"] in node_ids