graphitect 0.2.0__py3-none-any.whl

graphify/extractors/terraform.py

@@ -0,0 +1,181 @@
+"""Terraform extractor. Moved verbatim from graphify/extract.py."""
+from __future__ import annotations
+
+
+from pathlib import Path
+from graphify.extractors.base import _make_id
+
+
+_TF_META_HEADS = frozenset({"count", "each", "self", "path", "terraform"})
+
+def extract_terraform(path: Path) -> dict:
+    """Extract Terraform/HCL blocks and the references between them via tree-sitter.
+
+    Nodes: resources, data sources, modules, variables, outputs, providers, and
+    locals. Edges: `contains` (file -> block), `references` (block -> the blocks
+    it interpolates, e.g. `aws_instance.web` -> `var.region`), and `depends_on`
+    (explicit dependency edges).
+
+    Node IDs are scoped by the parent directory, not the file stem, because
+    Terraform resources are module(directory)-scoped: a resource defined in
+    main.tf is referenced from other .tf files in the same directory. Directory
+    scoping lets those cross-file references resolve when per-file extractions
+    are merged (stem scoping would split a definition from its references).
+    """
+    try:
+        import tree_sitter_hcl as tshcl
+        from tree_sitter import Language, Parser
+    except ImportError:
+        return {"nodes": [], "edges": [], "error": "tree_sitter_hcl not installed. Run: pip install tree-sitter-hcl"}
+
+    try:
+        language = Language(tshcl.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)}
+
+    str_path = str(path)
+    file_nid = _make_id(str_path)
+    scope = path.parent.name or "tf"
+
+    nodes: list[dict] = [{"id": file_nid, "label": path.name, "file_type": "code",
+                          "source_file": str_path, "source_location": None}]
+    edges: list[dict] = []
+    seen_ids: set[str] = {file_nid}
+    seen_edges: set[tuple[str, str, str]] = set()
+
+    def _read(n) -> str:
+        return source[n.start_byte:n.end_byte].decode("utf-8", errors="replace")
+
+    def _label_text(n) -> str:
+        return _read(n).strip().strip('"')
+
+    def _add_node(address: str, label: str, line: int) -> str:
+        nid = _make_id(scope, address)
+        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}"})
+            edges.append({"source": file_nid, "target": nid, "relation": "contains",
+                          "confidence": "EXTRACTED", "source_file": str_path,
+                          "source_location": f"L{line}", "weight": 1.0})
+        return nid
+
+    def _add_edge(src: str, address: str, relation: str, line: int) -> None:
+        tgt = _make_id(scope, address)
+        if src == tgt:
+            return
+        key = (src, tgt, relation)
+        if key in seen_edges:
+            return
+        seen_edges.add(key)
+        edges.append({"source": src, "target": tgt, "relation": relation,
+                      "confidence": "EXTRACTED", "source_file": str_path,
+                      "source_location": f"L{line}", "weight": 1.0})
+
+    def _block_parts(block) -> tuple:
+        btype = None
+        labels: list[str] = []
+        for c in block.children:
+            if c.type in ("block_start", "body", "block_end"):
+                break
+            if c.type == "identifier" and btype is None:
+                btype = _read(c)
+            elif c.type in ("string_lit", "identifier"):
+                labels.append(_label_text(c))
+        return btype, labels
+
+    def _ref_address(expr):
+        head = _read(expr)
+        parent = expr.parent
+        attrs: list[str] = []
+        if parent is not None:
+            seen_self = False
+            for c in parent.children:
+                if c.id == expr.id:
+                    seen_self = True
+                    continue
+                if seen_self and c.type == "get_attr":
+                    name = None
+                    for gc in c.children:
+                        if gc.type == "identifier":
+                            name = _read(gc)
+                            break
+                    if name is None:
+                        break
+                    attrs.append(name)
+                elif seen_self and c.type not in ("get_attr",):
+                    break
+        if head in _TF_META_HEADS or not head:
+            return None
+        if head == "var":
+            return f"var.{attrs[0]}" if attrs else None
+        if head == "local":
+            return f"local.{attrs[0]}" if attrs else None
+        if head == "module":
+            return f"module.{attrs[0]}" if attrs else None
+        if head == "data":
+            return f"data.{attrs[0]}.{attrs[1]}" if len(attrs) >= 2 else None
+        return f"{head}.{attrs[0]}" if attrs else None
+
+    def _collect_refs(node, owner_nid: str, relation: str) -> None:
+        rel = relation
+        if node.type == "attribute":
+            key_node = node.child_by_field_name("key") or (
+                node.children[0] if node.children else None
+            )
+            if key_node is not None and _read(key_node) == "depends_on":
+                rel = "depends_on"
+        if node.type == "variable_expr":
+            addr = _ref_address(node)
+            if addr:
+                _add_edge(owner_nid, addr, rel, node.start_point[0] + 1)
+        for c in node.children:
+            if c.is_named:
+                _collect_refs(c, owner_nid, rel)
+
+    def _body_of(block):
+        for c in block.children:
+            if c.type == "body":
+                return c
+        return None
+
+    body = next((c for c in root.children if c.type == "body"), root)
+    for block in body.children:
+        if block.type != "block":
+            continue
+        btype, labels = _block_parts(block)
+        line = block.start_point[0] + 1
+        blk_body = _body_of(block)
+        if btype == "resource" and len(labels) >= 2:
+            owner = _add_node(f"{labels[0]}.{labels[1]}", f"{labels[0]}.{labels[1]}", line)
+        elif btype == "data" and len(labels) >= 2:
+            owner = _add_node(f"data.{labels[0]}.{labels[1]}", f"data.{labels[0]}.{labels[1]}", line)
+        elif btype == "module" and labels:
+            owner = _add_node(f"module.{labels[0]}", f"module.{labels[0]}", line)
+        elif btype == "variable" and labels:
+            owner = _add_node(f"var.{labels[0]}", f"var.{labels[0]}", line)
+        elif btype == "output" and labels:
+            owner = _add_node(f"output.{labels[0]}", f"output.{labels[0]}", line)
+        elif btype == "provider" and labels:
+            owner = _add_node(f"provider.{labels[0]}", f"provider.{labels[0]}", line)
+        elif btype == "locals" and blk_body is not None:
+            for attr in blk_body.children:
+                if attr.type != "attribute":
+                    continue
+                key_node = attr.children[0] if attr.children else None
+                if key_node is None:
+                    continue
+                key = _read(key_node)
+                lnid = _add_node(f"local.{key}", f"local.{key}", attr.start_point[0] + 1)
+                _collect_refs(attr, lnid, "references")
+            continue
+        else:
+            continue
+        if blk_body is not None:
+            _collect_refs(blk_body, owner, "references")
+
+    return {"nodes": nodes, "edges": edges}

graphify/extractors/verilog.py

@@ -0,0 +1,329 @@
+"""Verilog extractor. Moved verbatim from graphify/extract.py."""
+from __future__ import annotations
+
+import re
+
+from pathlib import Path
+from graphify.extractors.base import _file_stem, _make_id, _read_text
+
+
+def _sv_first_identifier(node, source: bytes) -> str | None:
+    """First `simple_identifier` under node in pre-order, or None.
+
+    tree-sitter-verilog 1.0.3 nests declaration names a few levels deep instead
+    of exposing a `name` field. Scope the search to the right child node (e.g.
+    `function_identifier`) or this returns the return-type instead of the name.
+    """
+    if node is None:
+        return None
+    for child in node.children:
+        if child.type == "simple_identifier":
+            return _read_text(child, source)
+        found = _sv_first_identifier(child, source)
+        if found:
+            return found
+    return None
+
+def _sv_child(node, type_name: str) -> object | None:
+    if node is None:
+        return None
+    for child in node.children:
+        if child.type == type_name:
+            return child
+    return None
+
+_SV_BUILTIN_TYPES = frozenset({
+    "bit", "logic", "reg", "wire", "int", "integer", "shortint", "longint",
+    "byte", "time", "real", "shortreal", "void", "string", "type", "event",
+    "mailbox", "semaphore", "process", "chandle",
+})
+
+_SV_NON_TYPE_WORDS = frozenset({
+    "return", "if", "else", "for", "foreach", "while", "case", "begin", "end",
+    "function", "task", "class", "endclass", "endfunction", "endtask",
+})
+
+_SV_PARENS_INNER = r"(?:[^()]|\([^()]*\))*"
+
+_SV_PARENS = r"\(" + _SV_PARENS_INNER + r"\)"
+
+_SV_FUNC_RE = re.compile(
+    r"\bfunction\s+([A-Za-z_]\w*(?:\s*#\s*" + _SV_PARENS + r")?)\s+(\w+)\s*"
+    r"\((" + _SV_PARENS_INNER + r")\)\s*;",
+    re.MULTILINE,
+)
+
+_SV_PARAM_RE = re.compile(
+    r"\s*(?:input|output|inout|ref|const\s+ref)?\s*"
+    r"([A-Za-z_]\w*(?:\s*#\s*" + _SV_PARENS + r")?)\s+\w+"
+)
+
+def _sv_strip_comments(text: str) -> str:
+    text = re.sub(r"/\*.*?\*/", "", text, flags=re.DOTALL)
+    return re.sub(r"//.*", "", text)
+
+def _sv_split_type_list(text: str) -> list[str]:
+    parts: list[str] = []
+    depth = 0
+    start = 0
+    for idx, ch in enumerate(text):
+        if ch == "(":
+            depth += 1
+        elif ch == ")":
+            depth = max(0, depth - 1)
+        elif ch == "," and depth == 0:
+            item = text[start:idx].strip()
+            if item:
+                parts.append(item)
+            start = idx + 1
+    item = text[start:].strip()
+    if item:
+        parts.append(item)
+    return parts
+
+def _sv_collect_type_refs(type_text: str, generic: bool = False,
+                          skip: frozenset[str] = frozenset()) -> list[tuple[str, str]]:
+    refs: list[tuple[str, str]] = []
+    text = type_text.strip()
+    if not text:
+        return refs
+    head = re.match(r"([A-Za-z_]\w*)", text)
+    if head:
+        name = head.group(1)
+        # `skip` carries the enclosing class's `#(type T = ...)` parameters so
+        # they are not mistaken for referenced types.
+        if name not in _SV_BUILTIN_TYPES and name not in _SV_NON_TYPE_WORDS and name not in skip:
+            refs.append((name, "generic_arg" if generic else "type"))
+    params = re.search(r"#\s*\((" + _SV_PARENS_INNER + r")\)", text)
+    if params:
+        for arg in _sv_split_type_list(params.group(1)):
+            refs.extend(_sv_collect_type_refs(arg, generic=True, skip=skip))
+    return refs
+
+def _augment_systemverilog_semantics(
+    raw: str,
+    stem: str,
+    str_path: str,
+    file_nid: str,
+    nodes: list[dict],
+    edges: list[dict],
+    seen_ids: set[str],
+) -> None:
+    label_to_nid = {node["label"]: node["id"] for node in nodes}
+
+    def line_for(offset: int) -> int:
+        return raw.count("\n", 0, offset) + 1
+
+    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}",
+                          "confidence_score": 1.0})
+        label_to_nid[label] = nid
+
+    def ensure_type(label: str, line: int) -> str:
+        if label in label_to_nid:
+            return label_to_nid[label]
+        nid = _make_id(stem, label)
+        add_node(nid, label, line)
+        return nid
+
+    def add_edge(src: str, target_label: str, relation: str, line: int, context: str | None = None) -> None:
+        tgt = ensure_type(target_label, line)
+        edge = {"source": src, "target": tgt, "relation": relation,
+                "confidence": "EXTRACTED", "confidence_score": 1.0,
+                "source_file": str_path, "source_location": f"L{line}", "weight": 1.0}
+        if context:
+            edge["context"] = context
+        edges.append(edge)
+
+    text = _sv_strip_comments(raw)
+    # Consuming `endclass` (rather than a lookahead) makes each match own its
+    # terminator, so back-to-back or malformed classes cannot bleed bodies.
+    class_re = re.compile(
+        r"\b(?:(interface)\s+)?class\s+(\w+)([^;{]*)\s*;(.*?)\bendclass\b",
+        re.DOTALL,
+    )
+    for match in class_re.finditer(text):
+        class_name = match.group(2)
+        header = match.group(3) or ""
+        body = match.group(4) or ""
+        line = line_for(match.start())
+        # `#(type T = Payload)` declares `T` as a class type parameter, not a
+        # referenced type — collect these to skip below.
+        type_params = frozenset(re.findall(r"\btype\s+(\w+)", header))
+        class_nid = _make_id(stem, class_name)
+        add_node(class_nid, class_name, line)
+        edges.append({"source": file_nid, "target": class_nid, "relation": "defines",
+                      "confidence": "EXTRACTED", "confidence_score": 1.0,
+                      "source_file": str_path, "source_location": f"L{line}", "weight": 1.0})
+
+        ext = re.search(r"\bextends\s+(\w+)", header)
+        if ext:
+            add_edge(class_nid, ext.group(1), "inherits", line)
+        impl = re.search(r"\bimplements\s+([^;{]+)", header)
+        if impl:
+            for iface_name in _sv_split_type_list(impl.group(1)):
+                add_edge(class_nid, iface_name.split("#", 1)[0].strip(), "implements", line)
+
+        body_without_functions = re.sub(
+            r"\bfunction\b.*?\bendfunction\b",
+            lambda m: "\n" * m.group(0).count("\n"),
+            body,
+            flags=re.DOTALL,
+        )
+        # Optional leading class-property qualifiers (rand/local/protected/etc.)
+        # must be consumed: otherwise a qualified field like `rand Config x;`
+        # (three tokens) fails the `<type> <name>;` shape and its type reference
+        # is silently dropped.
+        for field in re.finditer(r"^\s*(?:(?:rand|randc|local|protected|static|const|automatic|var)\s+)*([A-Za-z_]\w*(?:\s*#\s*\([^;]+?\))?)\s+\w+\s*;", body_without_functions, re.MULTILINE):
+            # Count to the start of the type token (group 1), not the match
+            # start: `^\s*` consumes the leading newline(s), so field.start()
+            # would resolve to the class's line instead of the field's.
+            field_line = line + body_without_functions.count("\n", 0, field.start(1))
+            for ref_name, role in _sv_collect_type_refs(field.group(1), skip=type_params):
+                add_edge(class_nid, ref_name, "references", field_line, "generic_arg" if role == "generic_arg" else "field")
+
+        for fm in _SV_FUNC_RE.finditer(body):
+            return_type, func_name, params = fm.group(1), fm.group(2), fm.group(3)
+            func_line = line + body.count("\n", 0, fm.start())
+            func_nid = _make_id(class_nid, func_name)
+            add_node(func_nid, func_name, func_line)
+            edges.append({"source": class_nid, "target": func_nid, "relation": "method",
+                          "confidence": "EXTRACTED", "confidence_score": 1.0,
+                          "source_file": str_path, "source_location": f"L{func_line}", "weight": 1.0})
+            for ref_name, role in _sv_collect_type_refs(return_type, skip=type_params):
+                add_edge(func_nid, ref_name, "references", func_line, "generic_arg" if role == "generic_arg" else "return_type")
+            for param in _sv_split_type_list(params):
+                pm = _SV_PARAM_RE.match(param)
+                if not pm:
+                    continue
+                for ref_name, role in _sv_collect_type_refs(pm.group(1), skip=type_params):
+                    add_edge(func_nid, ref_name, "references", func_line, "generic_arg" if role == "generic_arg" else "parameter_type")
+
+def extract_verilog(path: Path) -> dict:
+    """Extract modules, functions, tasks, package imports, instantiations, and
+    SystemVerilog class semantics (inherits/implements edges, field/parameter/
+    return-type references) from .v/.sv files."""
+    try:
+        import tree_sitter_verilog as tsverilog
+        from tree_sitter import Language, Parser
+    except ImportError:
+        return {"nodes": [], "edges": [], "error": "tree_sitter_verilog not installed"}
+
+    try:
+        language = Language(tsverilog.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()
+
+    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}",
+                          "confidence_score": 1.0})
+
+    def add_edge(src: str, tgt: str, relation: str, line: int,
+                 confidence: str = "EXTRACTED", score: float = 1.0) -> None:
+        edges.append({"source": src, "target": tgt, "relation": relation,
+                      "confidence": confidence, "confidence_score": score,
+                      "source_file": str_path, "source_location": f"L{line}", "weight": 1.0})
+
+    file_nid = _make_id(str(path))
+    add_node(file_nid, path.name, 1)
+
+    def walk(node, module_nid: str | None = None) -> None:
+        t = node.type
+
+        # SystemVerilog class bodies are handled by _augment_systemverilog_semantics
+        # (regex over source text). Skip their subtrees so in-class methods are not
+        # double-emitted here — and with the wrong, return-type-derived name.
+        if t in ("class_declaration", "interface_class_declaration"):
+            return
+
+        if t == "module_declaration":
+            mod_name = _sv_first_identifier(_sv_child(node, "module_header"), source)
+            if mod_name:
+                line = node.start_point[0] + 1
+                nid = _make_id(stem, mod_name)
+                add_node(nid, mod_name, line)
+                add_edge(file_nid, nid, "defines", line)
+                for child in node.children:
+                    walk(child, nid)
+                return
+
+        # `function_prototype` only appears inside class/interface-class bodies
+        # (skipped above) and nests its name differently; it is intentionally not
+        # handled here.
+        elif t == "function_declaration":
+            fn_body = _sv_child(node, "function_body_declaration")
+            func_name = _sv_first_identifier(_sv_child(fn_body, "function_identifier"), source)
+            if func_name:
+                line = node.start_point[0] + 1
+                parent = module_nid or file_nid
+                nid = _make_id(parent, func_name)
+                add_node(nid, f"{func_name}()", line)
+                add_edge(parent, nid, "contains", line)
+
+        elif t == "task_declaration":
+            tk_body = _sv_child(node, "task_body_declaration")
+            task_name = _sv_first_identifier(_sv_child(tk_body, "task_identifier"), source)
+            if task_name:
+                line = node.start_point[0] + 1
+                parent = module_nid or file_nid
+                nid = _make_id(parent, task_name)
+                add_node(nid, task_name, line)
+                add_edge(parent, nid, "contains", line)
+
+        elif t == "package_import_declaration":
+            for child in node.children:
+                if child.type == "package_import_item":
+                    pkg_text = _read_text(child, source)
+                    pkg_name = pkg_text.split("::")[0].strip()
+                    if pkg_name:
+                        line = node.start_point[0] + 1
+                        tgt_nid = _make_id(pkg_name)
+                        add_node(tgt_nid, pkg_name, line)
+                        src_nid = module_nid or file_nid
+                        add_edge(src_nid, tgt_nid, "imports_from", line)
+
+        elif t in ("module_instantiation", "checker_instantiation"):
+            # `leaf u_leaf();` parses as checker_instantiation in 1.0.3;
+            # module_instantiation (when it occurs) exposes a `module_type` field.
+            # Both reduce to the first identifier under the node — the instantiated
+            # type, not the instance name (which appears later).
+            if module_nid:
+                type_node = node.child_by_field_name("module_type")
+                inst_type = (_read_text(type_node, source).strip() if type_node
+                             else _sv_first_identifier(node, source))
+                if inst_type:
+                    line = node.start_point[0] + 1
+                    tgt_nid = _make_id(inst_type)
+                    add_node(tgt_nid, inst_type, line)
+                    add_edge(module_nid, tgt_nid, "instantiates", line)
+
+        for child in node.children:
+            walk(child, module_nid)
+
+    walk(root)
+    _augment_systemverilog_semantics(
+        source.decode("utf-8", errors="replace"),
+        stem,
+        str_path,
+        file_nid,
+        nodes,
+        edges,
+        seen_ids,
+    )
+    return {"nodes": nodes, "edges": edges}