From 9d998cc81074efbfd052d72030d196e9d0b48ab8 Mon Sep 17 00:00:00 2001
From: Safi <safishamsi98@gmail.com>
Date: Tue, 7 Apr 2026 09:56:53 +0100
Subject: [PATCH] Add Zig and PowerShell language support

---
 README.md                 |   2 +-
 graphify/detect.py        |   2 +-
 graphify/extract.py       | 315 +++++++++++++++++++++++++++++++++++++-
 pyproject.toml            |   2 +
 tests/fixtures/sample.ps1 |  32 ++++
 tests/fixtures/sample.zig |  37 +++++
 tests/test_extract.py     |   2 +-
 7 files changed, 388 insertions(+), 4 deletions(-)
 create mode 100644 tests/fixtures/sample.ps1
 create mode 100644 tests/fixtures/sample.zig

diff --git a/README.md b/README.md
index ea94f8f1..dc800e8c 100644
--- a/README.md
+++ b/README.md
@@ -163,7 +163,7 @@ Works with any mix of file types:
 
 | Type | Extensions | Extraction |
 |------|-----------|------------|
-| Code | `.py .ts .js .go .rs .java .c .cpp .rb .cs .kt .scala .php .swift .lua` | AST via tree-sitter + call-graph + docstring/comment rationale |
+| Code | `.py .ts .js .go .rs .java .c .cpp .rb .cs .kt .scala .php .swift .lua .zig .ps1` | AST via tree-sitter + call-graph + docstring/comment rationale |
 | Docs | `.md .txt .rst` | Concepts + relationships + design rationale via Claude |
 | Papers | `.pdf` | Citation mining + concept extraction |
 | Images | `.png .jpg .webp .gif` | Claude vision - screenshots, diagrams, any language |
diff --git a/graphify/detect.py b/graphify/detect.py
index cb930b51..a033d677 100644
--- a/graphify/detect.py
+++ b/graphify/detect.py
@@ -17,7 +17,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', '.lua', '.toc'}
+CODE_EXTENSIONS = {'.py', '.ts', '.js', '.tsx', '.go', '.rs', '.java', '.cpp', '.cc', '.cxx', '.c', '.h', '.hpp', '.rb', '.swift', '.kt', '.kts', '.cs', '.scala', '.php', '.lua', '.toc', '.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 ccd976e6..3191e424 100644
--- a/graphify/extract.py
+++ b/graphify/extract.py
@@ -1482,6 +1482,317 @@ def extract_rust(path: Path) -> dict:
     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, Any]] = []
+
+    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:
+        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":
+                _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":
+            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
+
+        for child in node.children:
+            walk(child, parent_struct_nid)
+
+    walk(root)
+
+    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":
+            fn = node.child_by_field_name("function")
+            if fn:
+                callee = _read_text(fn, source).split(".")[-1]
+                tgt_nid = next((n["id"] for n in nodes if n["label"] in
+                                (f"{callee}()", f".{callee}()")), None)
+                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)
+
+    clean_edges = [e for e in edges if e["source"] in seen_ids and
+                   (e["target"] in seen_ids or e["relation"] == "imports_from")]
+    return {"nodes": nodes, "edges": clean_edges}
+
+
+# ── PowerShell ────────────────────────────────────────────────────────────────
+
+def extract_powershell(path: Path) -> dict:
+    """Extract functions, classes, methods, and using statements 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, Any]] = []
+
+    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 = 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):
+        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 = next((c for c in node.children if c.type == "function_name"), None)
+            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 = next((c for c in node.children if c.type == "simple_name"), None)
+            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 = next((c for c in node.children if c.type == "simple_name"), None)
+            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":
+            cmd_name_node = next((c for c in node.children if c.type == "command_name"), None)
+            if cmd_name_node:
+                cmd_text = _read_text(cmd_name_node, source).lower()
+                if cmd_text == "using":
+                    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))
+                    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]
+                        add_edge(file_nid, _make_id(module_name), "imports_from",
+                                 node.start_point[0] + 1)
+            return
+
+        for child in node.children:
+            walk(child, parent_class_nid)
+
+    walk(root)
+
+    label_to_nid = {n["label"].strip("()").lstrip(".").lower(): n["id"] for n in nodes}
+    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 = next((c for c in node.children if c.type == "command_name"), None)
+            if cmd_name_node:
+                cmd_text = _read_text(cmd_name_node, source)
+                if cmd_text.lower() not in _PS_SKIP:
+                    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)
+
+    clean_edges = [e for e in edges if e["source"] in seen_ids and
+                   (e["target"] in seen_ids or e["relation"] == "imports_from")]
+    return {"nodes": nodes, "edges": clean_edges}
+
+
 # ── Cross-file import resolution ──────────────────────────────────────────────
 
 def _resolve_cross_file_imports(
@@ -1667,6 +1978,8 @@ def extract(paths: list[Path]) -> dict:
         ".swift": extract_swift,
         ".lua": extract_lua,
         ".toc": extract_lua,
+        ".zig": extract_zig,
+        ".ps1": extract_powershell,
     }
 
     for path in paths:
@@ -1709,7 +2022,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", "*.swift",
-        "*.lua", "*.toc",
+        "*.lua", "*.toc", "*.zig", "*.ps1",
     )
     results: list[Path] = []
     for pattern in _EXTENSIONS:
diff --git a/pyproject.toml b/pyproject.toml
index 0a5b048d..f36e7ace 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,6 +28,8 @@ dependencies = [
     "tree-sitter-php",
     "tree-sitter-swift",
     "tree-sitter-lua",
+    "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 00000000..2cdb6aa7
--- /dev/null
+++ b/tests/fixtures/sample.ps1
@@ -0,0 +1,32 @@
+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)
+    Write-Output "Processing $Count items for $Name"
+}
+
+class DataProcessor {
+    [string]$Source
+
+    DataProcessor([string]$source) {
+        $this.Source = $source
+    }
+
+    [string] Transform([string]$input) {
+        return $input.ToUpper()
+    }
+
+    [void] Save([string]$path) {
+        Set-Content -Path $path -Value $this.Source
+    }
+}
diff --git a/tests/fixtures/sample.zig b/tests/fixtures/sample.zig
new file mode 100644
index 00000000..1a05c4af
--- /dev/null
+++ b/tests/fixtures/sample.zig
@@ -0,0 +1,37 @@
+const std = @import("std");
+const mem = @import("std").mem;
+
+const Point = struct {
+    x: f64,
+    y: f64,
+
+    pub fn distance(self: Point, other: Point) f64 {
+        const dx = self.x - other.x;
+        const dy = self.y - other.dy;
+        return std.math.sqrt(dx * dx + dy * dy);
+    }
+};
+
+const Color = enum {
+    red,
+    green,
+    blue,
+};
+
+const Shape = union(enum) {
+    circle: f64,
+    rect: Point,
+};
+
+pub fn add(a: i32, b: i32) i32 {
+    return a + b;
+}
+
+pub fn multiply(a: i32, b: i32) i32 {
+    return a * b;
+}
+
+pub fn main() void {
+    const result = add(1, 2);
+    _ = multiply(result, 3);
+}
diff --git a/tests/test_extract.py b/tests/test_extract.py
index e0afd6de..e95b273f 100644
--- a/tests/test_extract.py
+++ b/tests/test_extract.py
@@ -61,7 +61,7 @@ def test_collect_files_from_dir():
     supported = {".py", ".js", ".ts", ".tsx", ".go", ".rs",
                  ".java", ".c", ".cpp", ".cc", ".cxx", ".rb",
                  ".cs", ".kt", ".kts", ".scala", ".php", ".h", ".hpp",
-                 ".swift"}
+                 ".swift", ".lua", ".toc", ".zig", ".ps1"}
     assert all(f.suffix in supported for f in files)
     assert len(files) > 0