codegraph-gen 0.2.0__py3-none-any.whl

codegraph_gen/cluster.py

@@ -0,0 +1,116 @@
+import logging
+import networkx as nx
+from networkx.algorithms.community import louvain_communities
+
+logger = logging.getLogger(__name__)
+
+
+def detect_components(
+    G: nx.DiGraph,
+) -> tuple[dict[int, list[str]], dict[int, float], dict[int, str]]:
+    """
+    Detects logical components in the graph using modularity clustering.
+
+    Returns:
+        components: dict mapping component_id -> list of node_ids
+        cohesion_scores: dict mapping component_id -> cohesion density float
+        component_names: dict mapping component_id -> human friendly name
+    """
+    if G.number_of_nodes() == 0:
+        return {}, {}, {}
+
+    # Convert to undirected weighted graph for Louvain community detection
+    U = nx.Graph()
+    U.add_nodes_from(G.nodes)
+
+    for u, v, d in G.edges(data=True):
+        relation = d.get("relation")
+        if relation == "contains":
+            weight = 10.0
+        elif relation == "imports":
+            weight = 2.0
+        elif relation == "calls":
+            weight = 1.0
+        else:
+            weight = 1.0
+
+        if U.has_edge(u, v):
+            U[u][v]["weight"] += weight
+        else:
+            U.add_edge(u, v, weight=weight)
+
+    # Run Louvain community clustering with fixed seed for reproducibility
+    communities = list(louvain_communities(U, weight="weight", seed=42))
+
+    # Sort communities by size descending, breaking ties stably by sorted member IDs
+    communities.sort(key=lambda s: (-len(s), sorted(list(s))))
+
+    components = {}
+    cohesion_scores = {}
+    component_names = {}
+    raw_components = []
+
+    import os
+    from collections import Counter
+
+    for idx, member_set in enumerate(communities, start=1):
+        members = list(member_set)
+        components[idx] = members
+
+        # Calculate cohesion: density of the induced subgraph
+        subgraph = G.subgraph(members)
+        density = nx.density(subgraph)
+        cohesion_scores[idx] = round(density, 2)
+
+        # Name the component by its most central (highest degree) node
+        degrees = dict(G.degree(members))
+        if degrees:
+            # Sort by degree descending, and break ties alphabetically by node ID
+            sorted_nodes = sorted(degrees.keys(), key=lambda n: (-degrees[n], n))
+            most_central_node = sorted_nodes[0]
+            node_label = G.nodes[most_central_node].get("label", most_central_node)
+            # Remove trailing parens/extensions to make clean component name
+            clean_name = node_label.replace("()", "").split(".")[0]
+        else:
+            clean_name = f"Component {idx}"
+
+        # Find the longest common directory path
+        paths = []
+        for m in members:
+            sf = G.nodes[m].get("source_file")
+            if sf:
+                dir_path = os.path.dirname(sf)
+                if dir_path:
+                    paths.append(dir_path)
+
+        common_dir = ""
+        if paths:
+            try:
+                common_dir = os.path.commonpath(paths)
+                if common_dir in (".", "", "/"):
+                    common_dir = ""
+            except ValueError:
+                common_dir = ""
+
+        raw_components.append((idx, clean_name, common_dir))
+
+    # Count frequencies of candidate names to detect collisions
+    candidate_names = []
+    for idx, clean_name, common_dir in raw_components:
+        cand = common_dir if common_dir else clean_name
+        candidate_names.append(cand)
+    name_counts = Counter(candidate_names)
+
+    # Assign final unique component names
+    for idx, clean_name, common_dir in raw_components:
+        cand = common_dir if common_dir else clean_name
+        if name_counts[cand] == 1:
+            component_names[idx] = cand
+        else:
+            # Collision! Qualify the name to ensure uniqueness and clarity
+            if common_dir:
+                component_names[idx] = f"{common_dir} ({clean_name})"
+            else:
+                component_names[idx] = f"{clean_name} (Component {idx})"
+
+    return components, cohesion_scores, component_names

codegraph_gen/config.py

@@ -0,0 +1,76 @@
+import os
+from pathlib import Path
+from pydantic import BaseModel, Field
+from codegraph_gen.parser.base import ExtractionResult
+
+# Default exclusions for files and directories we want to ignore
+DEFAULT_EXCLUSIONS = {
+    ".git",
+    ".venv",
+    "venv",
+    "node_modules",
+    "third_party",
+    "dist",
+    "build",
+    ".build",
+    "__pycache__",
+    ".pytest_cache",
+    ".codegraph",
+    ".idea",
+    ".vscode",
+    "target",
+    "out",
+    "bin",
+    "obj",
+    "vendor",
+    "Pods",
+    "Carthage",
+    "DerivedData",
+    "build_output",
+    ".next",
+    ".nuxt",
+    ".cache",
+    "build_mac",
+    "build_ios",
+    "build_ios_sim",
+}
+
+
+# Mapping of supported languages to file extensions
+LANGUAGE_EXTENSIONS = {
+    "python": {".py"},
+    "javascript": {".js", ".mjs", ".cjs"},
+    "typescript": {".ts", ".tsx"},
+    "kotlin": {".kt", ".kts"},
+    "go": {".go"},
+    "rust": {".rs"},
+    "swift": {".swift"},
+    "c": {".c", ".h"},
+    "cpp": {".cpp", ".cc", ".cxx", ".hpp", ".hxx"},
+}
+
+ALL_EXTENSIONS = {ext for exts in LANGUAGE_EXTENSIONS.values() for ext in exts}
+
+
+class CacheEntry(BaseModel):
+    mtime: float
+    size: int
+    hash: str
+    result: ExtractionResult
+
+
+class CodegraphConfig(BaseModel):
+    """Configuration class for codegraph parsing and exporting."""
+
+    workspace_dir: Path
+    output_dir: Path = Field(default_factory=lambda: Path(".codegraph"))
+    exclusions: set[str] = Field(default_factory=lambda: DEFAULT_EXCLUSIONS)
+    languages: set[str] = Field(default_factory=lambda: set(LANGUAGE_EXTENSIONS.keys()))
+    max_workers: int = Field(default_factory=lambda: os.cpu_count() or 4)
+    use_cache: bool = Field(default=True)
+
+    @property
+    def absolute_output_dir(self) -> Path:
+        if self.output_dir.is_absolute():
+            return self.output_dir
+        return self.workspace_dir / self.output_dir

codegraph_gen/detect.py

@@ -0,0 +1,59 @@
+import logging
+from pathlib import Path
+from codegraph_gen.config import CodegraphConfig, LANGUAGE_EXTENSIONS
+
+logger = logging.getLogger(__name__)
+
+
+def discover_files(config: CodegraphConfig) -> list[tuple[Path, str]]:
+    """
+    Recursively discovers source files in the workspace directory.
+    Filters by allowed languages and ignores files/directories in exclusions.
+
+    Returns:
+        List of tuples: (absolute_file_path, language_name)
+    """
+    found_files = []
+    workspace = config.workspace_dir.resolve()
+
+    # Map extension -> language
+    ext_to_lang = {}
+    for lang in config.languages:
+        if lang in LANGUAGE_EXTENSIONS:
+            for ext in LANGUAGE_EXTENSIONS[lang]:
+                ext_to_lang[ext] = lang
+
+    # Normalize exclusions to lowercase for case-insensitive matching
+    exclusions_lower = {exc.lower() for exc in config.exclusions}
+
+    def is_ignored(path: Path) -> bool:
+        # Check if any part of the path is in exclusions_lower
+        try:
+            rel_parts = path.relative_to(workspace).parts
+        except ValueError:
+            # Not under workspace
+            return True
+
+        for part in rel_parts:
+            if part.lower() in exclusions_lower:
+                return True
+        return False
+
+    def scan_dir(directory: Path):
+        try:
+            for item in directory.iterdir():
+                if is_ignored(item):
+                    continue
+                if item.is_dir():
+                    scan_dir(item)
+                elif item.is_file():
+                    ext = item.suffix.lower()
+                    if ext in ext_to_lang:
+                        found_files.append((item.resolve(), ext_to_lang[ext]))
+        except PermissionError:
+            logger.warning(f"Permission denied: {directory}")
+        except Exception as e:
+            logger.error(f"Error scanning {directory}: {e}")
+
+    scan_dir(workspace)
+    return found_files

codegraph_gen/engine.py

@@ -0,0 +1,367 @@
+import logging
+import json
+import hashlib
+import concurrent.futures
+from enum import Enum
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Tuple
+import networkx as nx
+from pydantic import BaseModel, ConfigDict
+
+from codegraph_gen.config import CodegraphConfig, CacheEntry
+from codegraph_gen.parser.base import ExtractionResult
+from codegraph_gen.detect import discover_files
+from codegraph_gen.parser import get_parser
+from codegraph_gen.builder import build_graph
+from codegraph_gen.cluster import detect_components
+from codegraph_gen.analyzer import analyze_graph, AnalysisResult
+from codegraph_gen.renderer import (
+    MarkdownRenderer,
+    get_node_filename,
+    get_component_filename,
+)
+from codegraph_gen.writer import VaultWriter
+
+logger = logging.getLogger(__name__)
+
+
+def get_file_hash(path: Path) -> str:
+    """Computes MD5 hash of a file."""
+    hasher = hashlib.md5()
+    try:
+        with open(path, "rb") as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                hasher.update(chunk)
+    except Exception:
+        return ""
+    return hasher.hexdigest()
+
+
+def _parse_file_worker(
+    file_path: Path, lang: str, workspace_dir: Path
+) -> tuple[Path, Optional[ExtractionResult], Optional[str]]:
+    """Worker function for parallel file parsing."""
+    try:
+        from codegraph_gen.parser import get_parser
+
+        parser = get_parser(lang)
+        result = parser.parse_file(file_path, workspace_dir)
+        return file_path, result, None
+    except Exception as e:
+        import traceback
+
+        err_msg = f"{e}\n{traceback.format_exc()}"
+        return file_path, None, err_msg
+
+
+class PipelineStage(str, Enum):
+    DISCOVERING = "discovering"
+    PARSING = "parsing"
+    BUILDING = "building"
+    CLUSTERING = "clustering"
+    ANALYZING = "analyzing"
+    RENDERING = "rendering"
+    WRITING = "writing"
+    COMPLETED = "completed"
+
+
+class PipelineResult(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    graph: nx.DiGraph
+    files: List[Tuple[Path, str]]
+    components: Dict[int, List[str]]
+    cohesion_scores: Dict[int, float]
+    component_names: Dict[int, str]
+    analysis: AnalysisResult
+
+
+class CodegraphEngine:
+    def __init__(self, config: CodegraphConfig):
+        self.config = config
+        self.renderer = MarkdownRenderer(config.workspace_dir)
+        self.writer = VaultWriter()
+
+    def run_pipeline(
+        self,
+        progress_callback: Optional[
+            Callable[[PipelineStage, Any, int, int], None]
+        ] = None,
+    ) -> PipelineResult:
+        """
+        Runs the full codegraph generation pipeline.
+        Args:
+            progress_callback: A function taking (stage, current_item, index, total)
+        """
+        logger.info("Starting codegraph engine pipeline...")
+
+        # 1. Discover files
+        if progress_callback:
+            progress_callback(PipelineStage.DISCOVERING, None, 0, 0)
+        files = discover_files(self.config)
+        if not files:
+            logger.warning("No supported files found.")
+            if progress_callback:
+                progress_callback(PipelineStage.COMPLETED, None, 0, 0)
+            return PipelineResult(
+                graph=nx.DiGraph(),
+                files=[],
+                components={},
+                cohesion_scores={},
+                component_names={},
+                analysis=AnalysisResult(god_nodes=[], cycles=[], inter_comp_deps={}),
+            )
+
+        # 2. Parse files (with caching and optional parallel processing)
+        extractions = []
+        total_files = len(files)
+
+        cache_path = self.config.absolute_output_dir / "cache.json"
+        cache_entries = {}
+        if self.config.use_cache and cache_path.exists():
+            try:
+                with open(cache_path, "r", encoding="utf-8") as f:
+                    cache_data = json.load(f)
+                    for k, v in cache_data.items():
+                        cache_entries[k] = CacheEntry(**v)
+                logger.info(f"Loaded {len(cache_entries)} cache entries.")
+            except Exception as e:
+                logger.warning(f"Could not load cache: {e}")
+
+        files_to_parse = []
+        new_cache_entries = {}
+
+        for file_path, lang in files:
+            rel_path = str(file_path.relative_to(self.config.workspace_dir))
+            try:
+                stat = file_path.stat()
+                mtime = stat.st_mtime
+                size = stat.st_size
+                file_hash = get_file_hash(file_path)
+
+                # Check cache hit
+                if rel_path in cache_entries:
+                    entry = cache_entries[rel_path]
+                    if (
+                        entry.mtime == mtime
+                        and entry.size == size
+                        and entry.hash == file_hash
+                    ):
+                        extractions.append(entry.result)
+                        new_cache_entries[rel_path] = entry
+                        continue
+
+                # Cache miss
+                files_to_parse.append(
+                    (file_path, lang, rel_path, mtime, size, file_hash)
+                )
+            except Exception as e:
+                logger.error(f"Error accessing file metadata for {file_path}: {e}")
+                # Fallback to parsing without cache metadata
+                files_to_parse.append((file_path, lang, rel_path, 0.0, 0, ""))
+
+        num_hits = total_files - len(files_to_parse)
+        if num_hits > 0:
+            logger.info(
+                f"Cache hit: {num_hits} / {total_files} files loaded from cache."
+            )
+
+        if not files_to_parse:
+            if progress_callback:
+                progress_callback(PipelineStage.PARSING, None, total_files, total_files)
+        else:
+            max_workers = self.config.max_workers
+            if max_workers > 1 and len(files_to_parse) > 1:
+                logger.info(
+                    f"Parsing {len(files_to_parse)} files in parallel with {max_workers} workers..."
+                )
+                with concurrent.futures.ProcessPoolExecutor(
+                    max_workers=max_workers
+                ) as executor:
+                    futures = {
+                        executor.submit(
+                            _parse_file_worker,
+                            file_path,
+                            lang,
+                            self.config.workspace_dir,
+                        ): (file_path, rel_path, mtime, size, file_hash)
+                        for file_path, lang, rel_path, mtime, size, file_hash in files_to_parse
+                    }
+
+                    for idx, future in enumerate(
+                        concurrent.futures.as_completed(futures), start=1
+                    ):
+                        file_path, rel_path, mtime, size, file_hash = futures[future]
+                        progress_idx = num_hits + idx
+                        if progress_callback:
+                            progress_callback(
+                                PipelineStage.PARSING,
+                                file_path,
+                                progress_idx,
+                                total_files,
+                            )
+
+                        try:
+                            _, result, err_msg = future.result()
+                            if err_msg:
+                                logger.error(
+                                    f"Error parsing file {file_path} in worker: {err_msg}"
+                                )
+                            elif result:
+                                extractions.append(result)
+                                if file_hash:
+                                    new_cache_entries[rel_path] = CacheEntry(
+                                        mtime=mtime,
+                                        size=size,
+                                        hash=file_hash,
+                                        result=result,
+                                    )
+                        except Exception as e:
+                            logger.error(f"Failed to parse file {file_path}: {e}")
+            else:
+                logger.info(f"Parsing {len(files_to_parse)} files sequentially...")
+                for idx, (
+                    file_path,
+                    lang,
+                    rel_path,
+                    mtime,
+                    size,
+                    file_hash,
+                ) in enumerate(files_to_parse, start=1):
+                    progress_idx = num_hits + idx
+                    if progress_callback:
+                        progress_callback(
+                            PipelineStage.PARSING, file_path, progress_idx, total_files
+                        )
+                    try:
+                        parser = get_parser(lang)
+                        result = parser.parse_file(file_path, self.config.workspace_dir)
+                        extractions.append(result)
+                        if file_hash:
+                            new_cache_entries[rel_path] = CacheEntry(
+                                mtime=mtime, size=size, hash=file_hash, result=result
+                            )
+                    except Exception as e:
+                        logger.error(f"Error parsing file {file_path}: {e}")
+
+        # 3. Build graph
+        if progress_callback:
+            progress_callback(PipelineStage.BUILDING, None, 0, 0)
+        G = build_graph(extractions, self.config.workspace_dir)
+
+        # 4. Component clustering
+        if progress_callback:
+            progress_callback(PipelineStage.CLUSTERING, None, 0, 0)
+        components, cohesion_scores, component_names = detect_components(G)
+
+        # 5. Graph analysis
+        if progress_callback:
+            progress_callback(PipelineStage.ANALYZING, None, 0, 0)
+        analysis = analyze_graph(G, components)
+
+        # 6. Render pages in memory
+        if progress_callback:
+            progress_callback(PipelineStage.RENDERING, None, 0, 0)
+        node_component_map = {}
+        for cid, members in components.items():
+            comp_name = component_names.get(cid, f"Component {cid}")
+            for member in members:
+                node_component_map[member] = comp_name
+
+        rendered_nodes = {}
+        for nid, ndata in G.nodes(data=True):
+            fname = get_node_filename(nid)
+            content = self.renderer.render_node_page(nid, ndata, G, node_component_map)
+            rendered_nodes[fname] = content
+
+        rendered_components = {}
+        for cid, members in components.items():
+            comp_name = component_names[cid]
+            cohesion = cohesion_scores[cid]
+            fname = get_component_filename(comp_name)
+            content = self.renderer.render_component_page(
+                cid,
+                members,
+                G,
+                cohesion,
+                comp_name,
+                analysis.inter_comp_deps,
+                component_names,
+            )
+            rendered_components[fname] = content
+
+        # Check if README already has AI Insights and preserve it
+        ai_insights = None
+        readme_path = self.config.absolute_output_dir / "README.md"
+        if readme_path.exists():
+            try:
+                old_readme = readme_path.read_text(encoding="utf-8")
+                marker = None
+                for m in (
+                    "## AI Architectural Insights",
+                    "## AI 架构深度洞察 (AI Architectural Insights)",
+                    "## AI 架构深度洞察",
+                ):
+                    if m in old_readme:
+                        marker = m
+                        break
+                if marker:
+                    parts = old_readme.split(marker, 1)
+                    insights_text = parts[1].strip()
+                    if insights_text:
+                        ai_insights = insights_text
+            except Exception as e:
+                logger.warning(
+                    f"Could not read existing README.md to preserve AI insights: {e}"
+                )
+
+        readme_content = self.renderer.render_readme(
+            G,
+            components,
+            cohesion_scores,
+            component_names,
+            analysis,
+            ai_insights=ai_insights,
+        )
+
+        prompt_content = self.renderer.render_agent_prompt(
+            G, components, cohesion_scores, component_names, analysis
+        )
+
+        # 7. Write vault to disk
+        if progress_callback:
+            progress_callback(PipelineStage.WRITING, None, 0, 0)
+        self.writer.write_vault(
+            self.config.absolute_output_dir,
+            rendered_nodes,
+            rendered_components,
+            readme_content,
+            prompt_content,
+        )
+
+        # Write updated cache back to disk
+        if self.config.use_cache:
+            try:
+                self.config.absolute_output_dir.mkdir(parents=True, exist_ok=True)
+                with open(cache_path, "w", encoding="utf-8") as f:
+                    json.dump(
+                        {k: v.model_dump() for k, v in new_cache_entries.items()},
+                        f,
+                        indent=2,
+                    )
+                logger.info(f"Saved {len(new_cache_entries)} cache entries.")
+            except Exception as e:
+                logger.warning(f"Could not save cache: {e}")
+
+        if progress_callback:
+            progress_callback(PipelineStage.COMPLETED, None, 0, 0)
+
+        logger.info("Pipeline executed successfully.")
+        return PipelineResult(
+            graph=G,
+            files=files,
+            components=components,
+            cohesion_scores=cohesion_scores,
+            component_names=component_names,
+            analysis=analysis,
+        )

codegraph_gen/parser/__init__.py

@@ -0,0 +1,27 @@
+from codegraph_gen.parser.base import BaseParser
+from codegraph_gen.parser.python import PythonParser
+from codegraph_gen.parser.javascript import JavaScriptParser
+from codegraph_gen.parser.go import GoParser
+from codegraph_gen.parser.rust import RustParser
+from codegraph_gen.parser.swift import SwiftParser
+from codegraph_gen.parser.cpp import CParser, CppParser
+from codegraph_gen.parser.kotlin import KotlinParser
+
+PARSERS: dict[str, type[BaseParser]] = {
+    "python": PythonParser,
+    "javascript": JavaScriptParser,
+    "typescript": JavaScriptParser,  # uses same tree-sitter parser
+    "go": GoParser,
+    "rust": RustParser,
+    "swift": SwiftParser,
+    "c": CParser,
+    "cpp": CppParser,
+    "kotlin": KotlinParser,
+}
+
+
+def get_parser(language: str) -> BaseParser:
+    """Returns an instance of the parser for the given language."""
+    if language not in PARSERS:
+        raise ValueError(f"Unsupported language: {language}")
+    return PARSERS[language]()

codegraph_gen/parser/base.py

@@ -0,0 +1,38 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from pydantic import BaseModel
+
+
+class NodeSchema(BaseModel):
+    id: str  # Unique identifier, e.g. "relative_path::symbol_name"
+    label: str  # Human readable name, e.g. "my_function"
+    type: str  # 'file', 'class', 'function', 'method', 'struct', 'interface', 'trait', 'protocol'
+    source_file: str  # Path relative to workspace
+    line_start: int  # 1-indexed
+    line_end: int  # 1-indexed
+    signature: str  # Signature snippet
+    docstring: str = ""  # Docstring or comments
+    local_bindings: dict[
+        str, str
+    ] = {}  # Maps local variable/parameter name to its type name
+
+
+class EdgeSchema(BaseModel):
+    source: str  # Source node ID
+    target: str  # Target node ID
+    relation: str  # 'contains', 'imports', 'calls', 'inherits', 'implements'
+    import_map: dict[str, str] = {}  # Maps local name to original symbol name
+
+
+class ExtractionResult(BaseModel):
+    nodes: list[NodeSchema] = []
+    edges: list[EdgeSchema] = []
+
+
+class BaseParser(ABC):
+    """Abstract base class for all language-specific AST parsers."""
+
+    @abstractmethod
+    def parse_file(self, file_path: Path, workspace_dir: Path) -> ExtractionResult:
+        """Parses a file and extracts symbols (nodes) and relations (edges)."""
+        pass