coderay 1.0.0__py3-none-any.whl

coderay/graph/code_graph.py

@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from typing import Any
+
+import networkx as nx
+
+from coderay.chunking.registry import (
+    get_init_filenames,
+    get_resolution_suffixes,
+    get_supported_extensions,
+)
+from coderay.core.models import EdgeKind, GraphEdge, GraphNode, NodeKind
+
+logger = logging.getLogger(__name__)
+
+
+_KNOWN_EXTENSIONS: frozenset[str] = frozenset()
+_KNOWN_INIT_FILENAMES: frozenset[str] = frozenset()
+
+
+def _ensure_registry_cache() -> None:
+    """Lazily populate the cached extension and init filename sets."""
+    global _KNOWN_EXTENSIONS, _KNOWN_INIT_FILENAMES  # noqa: PLW0603
+    if not _KNOWN_EXTENSIONS:
+        _KNOWN_EXTENSIONS = frozenset(get_supported_extensions())
+        _KNOWN_INIT_FILENAMES = frozenset(get_init_filenames())
+
+
+def _file_path_to_module_names(file_path: str) -> list[str]:
+    """Derive possible module names from a file path."""
+    _ensure_registry_cache()
+
+    # Strip file extension using the registry
+    cleaned = file_path
+    for ext in sorted(_KNOWN_EXTENSIONS, key=len, reverse=True):
+        if cleaned.endswith(ext):
+            cleaned = cleaned[: -len(ext)]
+            break
+
+    parts = cleaned.replace("\\", "/").split("/")
+
+    # Strip common layout prefixes
+    if parts and parts[0] == "src":
+        parts = parts[1:]
+
+    # Strip init-style filenames (Python __init__, JS/TS index)
+    if parts and parts[-1] in _KNOWN_INIT_FILENAMES:
+        parts = parts[:-1]
+
+    if not parts:
+        return []
+
+    names: list[str] = []
+    for i in range(len(parts)):
+        suffix = parts[i:]
+        dotted = ".".join(suffix)
+        names.append(dotted)
+        slashed = "/".join(suffix)
+        if slashed != dotted:
+            names.append(slashed)
+    return names
+
+
+class CodeGraph:
+    """In-memory directed graph of code relationships."""
+
+    def __init__(self) -> None:
+        self._g: nx.DiGraph = nx.DiGraph()
+
+        # short name -> {full node IDs} — enables resolving bare names like
+        # "foo" to "src/utils.py::foo". Multiple IDs when the name is ambiguous
+        # (e.g. two files both define a function called "helper").
+        self._symbol_index: dict[str, set[str]] = defaultdict(set)
+
+        # dotted module name -> node ID — maps Python-style import paths
+        # (e.g. "core.models") to the MODULE node they refer to.
+        self._module_index: dict[str, str] = {}
+
+        # file path -> {node IDs} — all nodes belonging to a file, for O(k)
+        # file removal instead of scanning the entire graph.
+        self._file_index: dict[str, set[str]] = defaultdict(set)
+
+    def _index_node(self, node: GraphNode) -> None:
+        """Register a node in all secondary indexes."""
+        self._symbol_index[node.name].add(node.id)
+        self._file_index[node.file_path].add(node.id)
+        if node.kind == NodeKind.MODULE:
+            # Register all suffix variants so that "import models" and
+            # "import core.models" both resolve to the same MODULE node.
+            for mod_name in _file_path_to_module_names(node.file_path):
+                if mod_name not in self._module_index:
+                    self._module_index[mod_name] = node.id
+
+    def _unindex_node(self, node: GraphNode) -> None:
+        """Remove a node from all secondary indexes."""
+        sym_entries = self._symbol_index.get(node.name)
+        if sym_entries is not None:
+            sym_entries.discard(node.id)
+        file_entries = self._file_index.get(node.file_path)
+        if file_entries is not None:
+            file_entries.discard(node.id)
+        if node.kind == NodeKind.MODULE:
+            for mod_name in _file_path_to_module_names(node.file_path):
+                if self._module_index.get(mod_name) == node.id:
+                    del self._module_index[mod_name]
+
+    def add_node(self, node: GraphNode) -> None:
+        """Insert a node into the graph and update all secondary indexes."""
+        self._g.add_node(node.id, data=node)
+        self._index_node(node)
+
+    def add_edge(self, edge: GraphEdge) -> None:
+        """Insert a directed edge."""
+        self._g.add_edge(edge.source, edge.target, kind=edge.kind)
+
+    def add_nodes_and_edges(
+        self, nodes: list[GraphNode], edges: list[GraphEdge]
+    ) -> None:
+        """Bulk-insert nodes then edges (order matters — nodes first)."""
+        for n in nodes:
+            self.add_node(n)
+        for e in edges:
+            self.add_edge(e)
+
+    def remove_file(self, file_path: str) -> int:
+        """Remove all nodes belonging to a file and clean up every index.
+
+        Returns:
+            Number of nodes removed.
+        """
+        to_remove = self._file_index.pop(file_path, set())
+        for nid in to_remove:
+            node: GraphNode | None = self._g.nodes[nid].get("data")
+            if node:
+                self._unindex_node(node)
+            # NetworkX also removes all edges touching this node
+            self._g.remove_node(nid)
+        return len(to_remove)
+
+    def resolve_symbol(self, name: str, caller_file: str | None = None) -> str | None:
+        """Resolve a short/bare name to a fully-qualified node ID.
+
+        Returns:
+            Full node ID, or None if the name cannot be uniquely resolved.
+        """
+        # Already a full node ID (e.g. "src/a.py::foo") — fast path
+        if name in self._g and self._g.nodes[name].get("data") is not None:
+            return name
+        candidates = self._symbol_index.get(name, set())
+        if len(candidates) == 1:
+            return next(iter(candidates))
+        return None
+
+    def resolve_edges(self) -> int:
+        """Rewire phantom edge targets to real node IDs.
+
+        Returns:
+            Number of edges successfully resolved.
+        """
+        resolvable = {EdgeKind.CALLS, EdgeKind.INHERITS, EdgeKind.IMPORTS}
+        to_remove: list[tuple[str, str]] = []
+        to_add: list[tuple[str, str, dict]] = []
+
+        for u, v, data in self._g.edges(data=True):
+            kind = data.get("kind")
+            if kind not in resolvable:
+                continue
+            # Target is already a real node — nothing to resolve
+            if v in self._g and self._g.nodes[v].get("data") is not None:
+                continue
+
+            # Extract the caller's file path from its node ID
+            # ("src/a.py::MyClass.method" -> "src/a.py")
+            caller_file = u.split("::")[0] if "::" in u else u
+            resolved = self.resolve_symbol(v, caller_file=caller_file)
+            if not resolved and kind == EdgeKind.IMPORTS:
+                resolved = self._module_index.get(v)
+            if not resolved and kind == EdgeKind.IMPORTS:
+                resolved = self._resolve_path_target(v)
+            if resolved and resolved != v:
+                to_remove.append((u, v))
+                to_add.append((u, resolved, dict(data)))
+
+        # Apply changes in a second pass (safe to mutate now)
+        for u, v in to_remove:
+            if self._g.has_edge(u, v):
+                self._g.remove_edge(u, v)
+        for u, v, data in to_add:
+            self._g.add_edge(u, v, **data)
+
+        if to_add:
+            logger.info("Resolved %d edges via symbol/module index", len(to_add))
+        return len(to_add)
+
+    def _resolve_path_target(self, target: str) -> str | None:
+        """Try to match a path-style target to an existing MODULE node."""
+        if "/" not in target:
+            return None
+        for suffix in get_resolution_suffixes():
+            candidate = target + suffix
+            node_data = self._g.nodes.get(candidate, {})
+            if node_data and node_data.get("data") is not None:
+                return candidate
+        cleaned = target
+        if cleaned.startswith("src/"):
+            cleaned = cleaned[4:]
+        dotted = cleaned.replace("/", ".")
+        return self._module_index.get(dotted)
+
+    def get_node(self, node_id: str) -> GraphNode | None:
+        """Look up a node by its full ID. Returns None for phantoms or missing nodes."""
+        data = self._g.nodes.get(node_id)
+        if data:
+            return data.get("data")
+        return None
+
+    # ------------------------------------------------------------------
+    # Public query methods
+    # ------------------------------------------------------------------
+
+    def get_impact_radius(self, symbol: str, depth: int = 2) -> list[GraphNode]:
+        """Find all nodes that could be affected if ``symbol`` changes.
+
+        Args:
+            depth: Number of reverse-BFS hops. Higher values may return
+                very large sets.
+        """
+        resolved = self.resolve_symbol(symbol) or symbol
+        visited: set[str] = set()
+        frontier = {resolved}
+        for _ in range(depth):
+            next_frontier: set[str] = set()
+            for nid in frontier:
+                # predecessors = nodes that have an edge pointing TO nid
+                for pred in self._g.predecessors(nid):
+                    if pred not in visited:
+                        visited.add(pred)
+                        next_frontier.add(pred)
+            frontier = next_frontier
+        return [self.get_node(nid) for nid in visited if self.get_node(nid) is not None]
+
+    @property
+    def node_count(self) -> int:
+        """Total nodes in the graph (including phantoms)."""
+        return self._g.number_of_nodes()
+
+    @property
+    def edge_count(self) -> int:
+        """Total edges in the graph."""
+        return self._g.number_of_edges()
+
+    # ------------------------------------------------------------------
+    # Serialisation
+    #
+    # The graph is persisted as JSON (graph.json). Only real nodes are
+    # serialised — phantom nodes are recreated implicitly when edges
+    # referencing them are re-added. Secondary indexes are rebuilt by
+    # ``add_node`` during ``from_dict``.
+    # ------------------------------------------------------------------
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialise the graph to a JSON-compatible dict."""
+        nodes = []
+        for nid, data in self._g.nodes(data=True):
+            gn: GraphNode | None = data.get("data") if data else None
+            if gn:
+                nodes.append(
+                    {
+                        "id": gn.id,
+                        "kind": gn.kind.value,
+                        "file_path": gn.file_path,
+                        "start_line": gn.start_line,
+                        "end_line": gn.end_line,
+                        "name": gn.name,
+                        "qualified_name": gn.qualified_name,
+                    }
+                )
+        edges_list = []
+        for u, v, data in self._g.edges(data=True):
+            kind = data.get("kind", "")
+            kind_val = kind.value if hasattr(kind, "value") else str(kind)
+            edges_list.append({"source": u, "target": v, "kind": kind_val})
+        return {"nodes": nodes, "edges": edges_list}
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> CodeGraph:
+        """Deserialise a graph from a dict produced by ``to_dict``."""
+        graph = cls()
+        for nd in data.get("nodes", []):
+            graph.add_node(
+                GraphNode(
+                    id=nd["id"],
+                    kind=NodeKind(nd["kind"]),
+                    file_path=nd["file_path"],
+                    start_line=nd["start_line"],
+                    end_line=nd["end_line"],
+                    name=nd["name"],
+                    qualified_name=nd["qualified_name"],
+                )
+            )
+        for ed in data.get("edges", []):
+            graph.add_edge(
+                GraphEdge(
+                    source=ed["source"],
+                    target=ed["target"],
+                    kind=EdgeKind(ed["kind"]),
+                )
+            )
+        return graph

coderay/graph/extractor.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+import builtins
+import logging
+from typing import Any
+
+from coderay.chunking.registry import LanguageConfig, get_language_for_file
+from coderay.core.models import EdgeKind, GraphEdge, GraphNode, NodeKind
+
+logger = logging.getLogger(__name__)
+
+_PYTHON_BUILTINS: frozenset[str] = frozenset(
+    name for name in dir(builtins) if not name.startswith("_")
+)
+
+
+def build_callee_filter(config: dict[str, Any] | None = None) -> frozenset[str]:
+    """Build the callee exclusion set from builtins + user config.
+
+    Args:
+        config: Full application config dict. Only the ``graph`` section
+            is read. If None, defaults are used.
+
+    Returns:
+        Frozen set of callee names to exclude from CALLS edges.
+    """
+    graph_cfg = (config or {}).get("graph") or {}
+    extra_excludes = set(graph_cfg.get("exclude_callees") or [])
+    force_includes = set(graph_cfg.get("include_callees") or [])
+    return frozenset((_PYTHON_BUILTINS | extra_excludes) - force_includes)
+
+
+def _resolve_relative_import(source_file: str, relative_target: str) -> str | None:
+    """Resolve a Python relative import to a path-based target.
+
+    Args:
+        source_file: Path of the file containing the import.
+        relative_target: Dotted import string starting with one or more dots.
+
+    Returns:
+        Slash-separated path (no extension), or None if dots exceed
+        the directory depth.
+    """
+    dots = len(relative_target) - len(relative_target.lstrip("."))
+    rest = relative_target[dots:]
+
+    parts = source_file.replace("\\", "/").split("/")
+    dir_parts = parts[:-1]
+
+    levels_up = max(dots - 1, 0)
+    if levels_up > len(dir_parts):
+        return None
+    if levels_up > 0:
+        dir_parts = dir_parts[:-levels_up]
+
+    if rest:
+        dir_parts.extend(rest.split("."))
+
+    return "/".join(dir_parts) if dir_parts else None
+
+
+def _extract_callee_name(text: str) -> str:
+    """Extract the final method/function name from a call expression."""
+    cleaned = text
+    if cleaned.startswith("self."):
+        cleaned = cleaned[5:]
+    elif cleaned.startswith("this."):
+        cleaned = cleaned[5:]
+    parts = cleaned.split(".")
+    return parts[-1] if parts else cleaned
+
+
+class GraphExtractor:
+    """Extract graph nodes and edges from source files."""
+
+    def __init__(self, config: dict[str, Any] | None = None) -> None:
+        """Initialize the extractor with optional config overrides."""
+        self._excluded_callees = build_callee_filter(config)
+        self._source_bytes: bytes = b""
+        self._file_path: str = ""
+        self._module_id: str = ""
+        self._lang_cfg: LanguageConfig | None = None
+        self._nodes: list[GraphNode] = []
+        self._edges: list[GraphEdge] = []
+
+    def extract_from_file(
+        self,
+        file_path: str,
+        content: str,
+    ) -> tuple[list[GraphNode], list[GraphEdge]]:
+        """Parse a source file and extract all graph nodes and edges.
+
+        Returns:
+            Tuple of (nodes, edges). Returns ``([], [])`` if the language
+            is unsupported or parsing fails.
+        """
+        lang_cfg = get_language_for_file(file_path)
+        if lang_cfg is None:
+            return [], []
+
+        try:
+            parser = lang_cfg.get_parser()
+        except Exception:
+            return [], []
+
+        self._source_bytes = content.encode("utf-8")
+        self._file_path = file_path
+        self._module_id = file_path
+        self._lang_cfg = lang_cfg
+        self._nodes = []
+        self._edges = []
+
+        tree = parser.parse(self._source_bytes)
+
+        module_node = GraphNode(
+            id=self._module_id,
+            kind=NodeKind.MODULE,
+            file_path=file_path,
+            start_line=1,
+            end_line=tree.root_node.end_point[0] + 1,
+            name=file_path,
+            qualified_name=file_path,
+        )
+        self._nodes.append(module_node)
+
+        self._visit(tree.root_node, scope_stack=[])
+        return self._nodes, self._edges
+
+    # ------------------------------------------------------------------
+    # Tree traversal
+    # ------------------------------------------------------------------
+
+    def _visit(self, node, *, scope_stack: list[str]) -> None:
+        """Recursively walk the syntax tree, dispatching to type-specific handlers."""
+        ntype = node.type
+        lang_cfg = self._lang_cfg
+
+        if ntype in lang_cfg.import_types:
+            self._handle_import(node)
+        elif ntype in lang_cfg.function_scope_types:
+            self._handle_function_def(node, scope_stack=scope_stack)
+            return
+        elif ntype in lang_cfg.class_scope_types:
+            self._handle_class_def(node, scope_stack=scope_stack)
+            return
+        elif ntype in lang_cfg.call_types:
+            self._handle_call(node, scope_stack=scope_stack)
+
+        for child in node.children:
+            self._visit(child, scope_stack=scope_stack)
+
+    # ------------------------------------------------------------------
+    # Node-type handlers
+    # ------------------------------------------------------------------
+
+    def _handle_import(self, node) -> None:
+        """Create IMPORTS edges for an import statement."""
+        is_from_import = node.type == "import_from_statement"
+        found_module = False
+        for child in node.children:
+            if child.type in ("dotted_name", "relative_import"):
+                if is_from_import and found_module:
+                    continue
+                target = self._text(child)
+                if target:
+                    if child.type == "relative_import" and target.startswith("."):
+                        resolved = _resolve_relative_import(self._module_id, target)
+                        if resolved:
+                            target = resolved
+                    self._edges.append(
+                        GraphEdge(
+                            source=self._module_id,
+                            target=target,
+                            kind=EdgeKind.IMPORTS,
+                        )
+                    )
+                    if is_from_import:
+                        found_module = True
+            elif child.type == "string":
+                target = self._text(child).strip("'\"")
+                if target:
+                    self._edges.append(
+                        GraphEdge(
+                            source=self._module_id,
+                            target=target,
+                            kind=EdgeKind.IMPORTS,
+                        )
+                    )
+            elif child.type == "interpreted_string_literal":
+                target = self._text(child).strip('"')
+                if target:
+                    self._edges.append(
+                        GraphEdge(
+                            source=self._module_id,
+                            target=target,
+                            kind=EdgeKind.IMPORTS,
+                        )
+                    )
+
+    def _handle_function_def(self, node, *, scope_stack: list[str]) -> None:
+        """Create a FUNCTION node and DEFINES edge, then recurse into the body."""
+        name = self._get_identifier(node)
+        if not name:
+            return
+        qualified = ".".join([*scope_stack, name])
+        node_id = f"{self._file_path}::{qualified}"
+        self._nodes.append(
+            GraphNode(
+                id=node_id,
+                kind=NodeKind.FUNCTION,
+                file_path=self._file_path,
+                start_line=node.start_point[0] + 1,
+                end_line=node.end_point[0] + 1,
+                name=name,
+                qualified_name=qualified,
+            )
+        )
+        self._edges.append(
+            GraphEdge(
+                source=self._module_id,
+                target=node_id,
+                kind=EdgeKind.DEFINES,
+            )
+        )
+        new_scope = [*scope_stack, name]
+        for child in node.children:
+            self._visit(child, scope_stack=new_scope)
+
+    def _handle_class_def(self, node, *, scope_stack: list[str]) -> None:
+        """Create a CLASS node, DEFINES + INHERITS edges, then recurse."""
+        name = self._get_identifier(node)
+        if not name:
+            return
+        qualified = ".".join([*scope_stack, name])
+        node_id = f"{self._file_path}::{qualified}"
+        self._nodes.append(
+            GraphNode(
+                id=node_id,
+                kind=NodeKind.CLASS,
+                file_path=self._file_path,
+                start_line=node.start_point[0] + 1,
+                end_line=node.end_point[0] + 1,
+                name=name,
+                qualified_name=qualified,
+            )
+        )
+        self._edges.append(
+            GraphEdge(
+                source=self._module_id,
+                target=node_id,
+                kind=EdgeKind.DEFINES,
+            )
+        )
+        for child in node.children:
+            if child.type in ("argument_list", "superclass", "extends_clause"):
+                for arg in child.children:
+                    if arg.type in (
+                        "identifier",
+                        "dotted_name",
+                        "attribute",
+                        "type_identifier",
+                    ):
+                        if base_name := self._text(arg):
+                            self._edges.append(
+                                GraphEdge(
+                                    source=node_id,
+                                    target=base_name,
+                                    kind=EdgeKind.INHERITS,
+                                )
+                            )
+        new_scope = [*scope_stack, name]
+        for child in node.children:
+            self._visit(child, scope_stack=new_scope)
+
+    def _handle_call(self, node, *, scope_stack: list[str]) -> None:
+        """Create a CALLS edge from the enclosing scope to the callee."""
+        caller_id = (
+            f"{self._file_path}::{'.'.join(scope_stack)}"
+            if scope_stack
+            else self._module_id
+        )
+        first_child = node.children[0] if node.children else None
+        if first_child is None:
+            return
+        raw_callee = self._text(first_child)
+        if not raw_callee:
+            return
+        callee_name = _extract_callee_name(raw_callee)
+        if callee_name and callee_name not in self._excluded_callees:
+            self._edges.append(
+                GraphEdge(
+                    source=caller_id,
+                    target=callee_name,
+                    kind=EdgeKind.CALLS,
+                )
+            )
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _get_identifier(self, node) -> str:
+        """Return the identifier name from a definition node."""
+        for child in node.children:
+            if child.type in ("identifier", "type_identifier", "field_identifier"):
+                return self._source_bytes[child.start_byte : child.end_byte].decode(
+                    "utf-8", errors="replace"
+                )
+        return ""
+
+    def _text(self, node) -> str:
+        """Decode the raw source text spanned by a syntax tree node."""
+        return self._source_bytes[node.start_byte : node.end_byte].decode(
+            "utf-8", errors="replace"
+        )