graphitect 0.2.0__py3-none-any.whl

graphify/extractors/fortran.py

@@ -0,0 +1,311 @@
+"""Fortran extractor. Moved verbatim from graphify/extract.py."""
+from __future__ import annotations
+
+
+from pathlib import Path
+from graphify.extractors.base import _file_stem, _make_id, _read_text
+
+
+_FORTRAN_CPP_EXTS = {".F", ".F90", ".F95", ".F03", ".F08"}
+
+def _cpp_preprocess(path: Path) -> bytes:
+    """Run cpp -w -P on a capital-F Fortran file and return preprocessed bytes.
+
+    Falls back to raw file bytes if cpp is not available. Capital-F extensions
+    conventionally require C preprocessor expansion (#ifdef MPI, #define REAL8, etc.)
+    before parsing.
+
+    Security (F-007): we pass `-nostdinc` and `-I /dev/null` so a malicious
+    source file containing `#include "/home/victim/.ssh/id_rsa"` (or any other
+    include directive) cannot inline arbitrary host files into the output that
+    we then ship to an LLM. Without these flags `cpp` happily resolves any
+    relative or absolute include path it can read, which is a corpus-side
+    file-exfiltration vector.
+    """
+    import shutil
+    import subprocess
+    if not shutil.which("cpp"):
+        return path.read_bytes()
+    try:
+        # Pass an absolute path so a corpus file named like "-I/etc/x.F90" cannot
+        # be parsed by cpp as an option (cpp does not accept a "--" end-of-options
+        # terminator). What matters is that an absolute path cannot begin with
+        # "-" — not that it begins with "/", which only holds on POSIX (a Windows
+        # absolute path starts with a drive letter, and is equally safe).
+        result = subprocess.run(
+            ["cpp", "-w", "-P", "-nostdinc", "-I", "/dev/null", str(path.resolve())],
+            capture_output=True,
+            timeout=30,
+        )
+        if result.returncode == 0 and result.stdout:
+            return result.stdout
+    except Exception:
+        pass
+    return path.read_bytes()
+
+def extract_fortran(path: Path) -> dict:
+    """Extract programs, modules, subroutines, functions, use statements, and calls from Fortran files.
+
+    Capital-F extensions (.F, .F90, etc.) are run through the C preprocessor before
+    parsing so #ifdef/#define macros are resolved.
+    """
+    try:
+        import tree_sitter_fortran as tsfortran
+        from tree_sitter import Language, Parser
+    except ImportError:
+        return {"nodes": [], "edges": [], "error": "tree-sitter-fortran not installed"}
+
+    try:
+        language = Language(tsfortran.language())
+        parser = Parser(language)
+        source = _cpp_preprocess(path) if path.suffix in _FORTRAN_CPP_EXTS else 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()
+    scope_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,
+                 context: str | None = None) -> None:
+        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)
+
+    def _fortran_name(stmt_node) -> str | None:
+        """Extract name from a *_statement node. Fortran is case-insensitive; lowercase."""
+        for child in stmt_node.children:
+            if child.type in ("name", "identifier"):
+                return _read_text(child, source).lower()
+        return None
+
+    def ensure_named_node(name: str, line: int) -> str:
+        nid = _make_id(stem, name)
+        if nid in seen_ids:
+            return nid
+        nid = _make_id(name)
+        if nid not in seen_ids:
+            # The name isn't defined in this file, so this is a cross-file reference
+            # (e.g. a `Thing` type annotation imported from another module). Emit a
+            # SOURCELESS stub — like the inheritance-base path below — so the
+            # corpus-level rewire can collapse it onto the real definition. A sourced
+            # stub here makes _disambiguate_colliding_node_ids bake the referencing
+            # file's path (with extension) into the id and blocks the rewire, which is
+            # the phantom-duplicate-node bug (#1402).
+            seen_ids.add(nid)
+            nodes.append({
+                "id": nid,
+                "label": name,
+                "file_type": "code",
+                "source_file": "",
+                "source_location": "",
+                "origin_file": str_path,
+            })
+        return nid
+
+    def emit_signature_refs(scope_node, fn_nid: str, is_function: bool) -> None:
+        """Emit references[parameter_type] / references[return_type] edges for
+        a subroutine/function based on its variable_declaration siblings."""
+        stmt_type = "function_statement" if is_function else "subroutine_statement"
+        stmt = next((c for c in scope_node.children if c.type == stmt_type), None)
+        if stmt is None:
+            return
+        param_names: set[str] = set()
+        params_node = next((c for c in stmt.children if c.type == "parameters"), None)
+        if params_node is not None:
+            for c in params_node.children:
+                if c.type == "identifier":
+                    param_names.add(_read_text(c, source).lower())
+        result_name: str | None = None
+        if is_function:
+            result_node = next((c for c in stmt.children if c.type == "function_result"), None)
+            if result_node is not None:
+                res_id = next((c for c in result_node.children if c.type == "identifier"), None)
+                if res_id is not None:
+                    result_name = _read_text(res_id, source).lower()
+            else:
+                # implicit result variable: same name as the function
+                result_name = _fortran_name(stmt)
+        for child in scope_node.children:
+            if child.type != "variable_declaration":
+                continue
+            derived = next((c for c in child.children if c.type == "derived_type"), None)
+            if derived is None:
+                continue
+            type_name_node = next((c for c in derived.children if c.type == "type_name"), None)
+            if type_name_node is None:
+                continue
+            type_name = _read_text(type_name_node, source).lower()
+            for var in child.children:
+                if var.type != "identifier":
+                    continue
+                var_name = _read_text(var, source).lower()
+                var_line = var.start_point[0] + 1
+                if var_name in param_names:
+                    tgt = ensure_named_node(type_name, var_line)
+                    if tgt != fn_nid:
+                        add_edge(fn_nid, tgt, "references", var_line, context="parameter_type")
+                elif is_function and var_name == result_name:
+                    tgt = ensure_named_node(type_name, var_line)
+                    if tgt != fn_nid:
+                        add_edge(fn_nid, tgt, "references", var_line, context="return_type")
+
+    def walk_calls(node, scope_nid: str) -> None:
+        if node is None:
+            return
+        t = node.type
+        if t in ("subroutine", "function", "module", "program", "internal_procedures"):
+            return
+        # call FOO(args) — tree-sitter-fortran uses subroutine_call
+        if t == "subroutine_call":
+            name_node = next((c for c in node.children if c.type == "identifier"), None)
+            if name_node:
+                callee = _read_text(name_node, source).lower()
+                target_nid = _make_id(stem, callee)
+                add_edge(scope_nid, target_nid, "calls", node.start_point[0] + 1,
+                         confidence="EXTRACTED", context="call")
+        # x = compute(args) — function invocations are `call_expression`, which
+        # shares Fortran's `name(...)` syntax with array indexing. Only emit a
+        # call edge when the callee resolves to a procedure defined in this file
+        # (an array variable produces no matching node), so array accesses can't
+        # fabricate spurious `calls` edges.
+        elif t == "call_expression":
+            name_node = next((c for c in node.children if c.type == "identifier"), None)
+            if name_node:
+                callee = _read_text(name_node, source).lower()
+                target_nid = _make_id(stem, callee)
+                if target_nid in seen_ids and target_nid != scope_nid:
+                    add_edge(scope_nid, target_nid, "calls", node.start_point[0] + 1,
+                             confidence="EXTRACTED", context="call")
+        for child in node.children:
+            walk_calls(child, scope_nid)
+
+    def walk(node, scope_nid: str) -> None:
+        t = node.type
+
+        if t == "program":
+            stmt = next((c for c in node.children if c.type == "program_statement"), None)
+            name = _fortran_name(stmt) if stmt else None
+            if name:
+                nid = _make_id(stem, name)
+                line = node.start_point[0] + 1
+                add_node(nid, name, line)
+                add_edge(file_nid, nid, "defines", line)
+                scope_bodies.append((nid, node))
+                for child in node.children:
+                    walk(child, nid)
+            return
+
+        if t == "module":
+            stmt = next((c for c in node.children if c.type == "module_statement"), None)
+            name = _fortran_name(stmt) if stmt else None
+            if name:
+                nid = _make_id(stem, name)
+                line = node.start_point[0] + 1
+                add_node(nid, name, line)
+                add_edge(file_nid, nid, "defines", line)
+                for child in node.children:
+                    walk(child, nid)
+            return
+
+        # subroutines/functions inside a module live under internal_procedures
+        if t == "internal_procedures":
+            for child in node.children:
+                walk(child, scope_nid)
+            return
+
+        if t == "derived_type_definition":
+            stmt = next((c for c in node.children if c.type == "derived_type_statement"), None)
+            if stmt is not None:
+                name_node = next((c for c in stmt.children if c.type == "type_name"), None)
+                if name_node is not None:
+                    type_name = _read_text(name_node, source).lower()
+                    type_nid = _make_id(stem, type_name)
+                    line = node.start_point[0] + 1
+                    add_node(type_nid, type_name, line)
+                    add_edge(scope_nid, type_nid, "defines", line)
+            return
+
+        if t == "subroutine":
+            stmt = next((c for c in node.children if c.type == "subroutine_statement"), None)
+            name = _fortran_name(stmt) if stmt else None
+            if name:
+                nid = _make_id(stem, name)
+                line = node.start_point[0] + 1
+                add_node(nid, f"{name}()", line)
+                add_edge(scope_nid, nid, "defines", line)
+                scope_bodies.append((nid, node))
+                emit_signature_refs(node, nid, is_function=False)
+                for child in node.children:
+                    walk(child, nid)
+            return
+
+        if t == "function":
+            stmt = next((c for c in node.children if c.type == "function_statement"), None)
+            name = _fortran_name(stmt) if stmt else None
+            if name:
+                nid = _make_id(stem, name)
+                line = node.start_point[0] + 1
+                add_node(nid, f"{name}()", line)
+                add_edge(scope_nid, nid, "defines", line)
+                scope_bodies.append((nid, node))
+                emit_signature_refs(node, nid, is_function=True)
+                for child in node.children:
+                    walk(child, nid)
+            return
+
+        if t == "use_statement":
+            line = node.start_point[0] + 1
+            # tree-sitter-fortran uses module_name node for the used module
+            name_node = next((c for c in node.children if c.type in ("module_name", "name", "identifier")), None)
+            if name_node:
+                mod_name = _read_text(name_node, source).lower()
+                imp_nid = _make_id(mod_name)
+                add_node(imp_nid, mod_name, line)
+                add_edge(scope_nid, imp_nid, "imports", line, context="use")
+            return
+
+        for child in node.children:
+            walk(child, scope_nid)
+
+    walk(root, file_nid)
+
+    _stmt_headers = {
+        "subroutine_statement", "function_statement",
+        "program_statement", "module_statement",
+    }
+    for scope_nid, body_node in scope_bodies:
+        for child in body_node.children:
+            if child.type not in _stmt_headers:
+                walk_calls(child, scope_nid)
+
+    return {"nodes": nodes, "edges": edges}