diff --git a/code2llm/core/models.py b/code2llm/core/models.py index 8b5a2ad..1075faf 100644 --- a/code2llm/core/models.py +++ b/code2llm/core/models.py @@ -1,207 +1,227 @@ -from dataclasses import dataclass, field, asdict -from typing import List, Dict, Set, Optional, Any - - -class BaseModel: - """Base class for models with automated serialization.""" - - def to_dict(self, compact: bool = True) -> dict: - """Convert to dictionary using dataclasses.asdict with filtering.""" - data = asdict(self) - if compact: - return self._filter_compact(data) - return data - - def _filter_compact(self, data: Any) -> Any: - """Recursively filter out None and empty collections if compact.""" - if isinstance(data, dict): - return { - k: self._filter_compact(v) - for k, v in data.items() - if v is not None - and (not isinstance(v, (list, dict, set)) or len(v) > 0) - } - elif isinstance(data, (list, tuple, set)): - return [self._filter_compact(v) for v in data] - return data - - -@dataclass -class FlowNode(BaseModel): - """Represents a node in the control flow graph.""" - - id: str - type: str # FUNC, CALL, IF, FOR, WHILE, ASSIGN, RETURN, ENTRY, EXIT - label: str - function: Optional[str] = None - file: Optional[str] = None - line: Optional[int] = None - column: Optional[int] = None - conditions: List[str] = field(default_factory=list) - data_flow: List[str] = field(default_factory=list) - metadata: Dict[str, Any] = field(default_factory=dict) - - -@dataclass -class FlowEdge(BaseModel): - """Represents an edge in the control flow graph.""" - - source: str - target: str - edge_type: str = "control" # control, data, call - label: Optional[str] = None - conditions: List[str] = field(default_factory=list) - - -@dataclass -class FunctionInfo(BaseModel): - """Information about a function/method.""" - - name: str - qualified_name: str - file: str - line: int - column: int = 0 - module: str = "" - class_name: Optional[str] = None - is_method: bool = False - is_private: bool = False - is_property: bool = False - docstring: Optional[str] = None - args: List[str] = field(default_factory=list) - returns: Optional[str] = None - decorators: List[str] = field(default_factory=list) - - # CFG info - cfg_entry: Optional[str] = None - cfg_exit: Optional[str] = None - cfg_nodes: List[str] = field(default_factory=list) - calls: List[str] = field(default_factory=list) - called_by: List[str] = field(default_factory=list) - - # Advanced metrics (Sprint 3) - complexity: Dict[str, Any] = field(default_factory=dict) # Cyclomatic, Cognitive - centrality: float = 0.0 # Betweenness Centrality - reachability: str = "unknown" # reachable, unreachable, unknown - - -@dataclass -class ClassInfo(BaseModel): - """Information about a class.""" - - name: str - qualified_name: str - file: str - line: int - module: str = "" - bases: List[str] = field(default_factory=list) - methods: List[str] = field(default_factory=list) - docstring: Optional[str] = None - is_state_machine: bool = False - - -@dataclass -class ModuleInfo(BaseModel): - """Information about a module/package.""" - - name: str - file: str - is_package: bool = False - source_kind: str = "source" - line_count: int = 0 - imports: List[str] = field(default_factory=list) - functions: List[str] = field(default_factory=list) - classes: List[str] = field(default_factory=list) - - -@dataclass -class Pattern(BaseModel): - """Detected behavioral pattern.""" - - name: str - type: str # recursion, state_machine, factory, singleton, strategy, loop - confidence: float # 0.0 to 1.0 - functions: List[str] = field(default_factory=list) - entry_points: List[str] = field(default_factory=list) - exit_points: List[str] = field(default_factory=list) - metadata: Dict[str, Any] = field(default_factory=dict) - - -@dataclass -class CodeSmell(BaseModel): - """Represents a detected code smell.""" - - name: str - type: str # god_function, feature_envy, etc. - file: str - line: int - severity: float # 0.0 to 1.0 - description: str - context: Dict[str, Any] = field(default_factory=dict) - - -@dataclass -class Mutation(BaseModel): - """Represents a mutation of a variable/object.""" - - variable: str - file: str - line: int - type: str # assign, aug_assign, method_call - scope: str - context: str - - -@dataclass -class DataFlow(BaseModel): - """Represents data flow for a variable.""" - - variable: str - dependencies: Set[str] = field(default_factory=set) - metadata: Dict[str, Any] = field(default_factory=dict) - - -@dataclass -class AnalysisResult(BaseModel): - """Complete analysis result for a project.""" - - project_path: str = "" - analysis_mode: str = "static" - stats: Dict[str, int] = field(default_factory=dict) - - # Graph data - nodes: Dict[str, FlowNode] = field(default_factory=dict) - edges: List[FlowEdge] = field(default_factory=list) - - # Code structure - modules: Dict[str, ModuleInfo] = field(default_factory=dict) - classes: Dict[str, ClassInfo] = field(default_factory=dict) - functions: Dict[str, FunctionInfo] = field(default_factory=dict) - - # Analysis results - patterns: List[Pattern] = field(default_factory=list) - call_graph: Dict[str, List[str]] = field(default_factory=list) - entry_points: List[str] = field(default_factory=list) - data_flows: Dict[str, DataFlow] = field(default_factory=dict) - - # Refactoring data - metrics: Dict[str, Dict[str, Any]] = field(default_factory=dict) - smells: List[CodeSmell] = field(default_factory=list) - coupling: Dict[str, Any] = field(default_factory=dict) - mutations: List[Mutation] = field(default_factory=list) - - def get_function_count(self) -> int: - """Get total function count.""" - return len(self.functions) - - def get_class_count(self) -> int: - """Get total class count.""" - return len(self.classes) - - def get_node_count(self) -> int: - """Get total CFG node count.""" - return len(self.nodes) - - def get_edge_count(self) -> int: - """Get total edge count.""" - return len(self.edges) +"""Data models for code analysis. + +This module re-exports models from the independent `code2graph` package +if installed, or falls back to the embedded definitions. +""" +try: + from code2graph.models import ( + BaseModel, + FlowNode, + FlowEdge, + FunctionInfo, + ClassInfo, + ModuleInfo, + Pattern, + CodeSmell, + Mutation, + DataFlow, + AnalysisResult, + ) +except ImportError: + from dataclasses import dataclass, field, asdict + from typing import List, Dict, Set, Optional, Any + + + class BaseModel: + """Base class for models with automated serialization.""" + + def to_dict(self, compact: bool = True) -> dict: + """Convert to dictionary using dataclasses.asdict with filtering.""" + data = asdict(self) + if compact: + return self._filter_compact(data) + return data + + def _filter_compact(self, data: Any) -> Any: + """Recursively filter out None and empty collections if compact.""" + if isinstance(data, dict): + return { + k: self._filter_compact(v) + for k, v in data.items() + if v is not None + and (not isinstance(v, (list, dict, set)) or len(v) > 0) + } + elif isinstance(data, (list, tuple, set)): + return [self._filter_compact(v) for v in data] + return data + + + @dataclass + class FlowNode(BaseModel): + """Represents a node in the control flow graph.""" + + id: str + type: str # FUNC, CALL, IF, FOR, WHILE, ASSIGN, RETURN, ENTRY, EXIT + label: str + function: Optional[str] = None + file: Optional[str] = None + line: Optional[int] = None + column: Optional[int] = None + conditions: List[str] = field(default_factory=list) + data_flow: List[str] = field(default_factory=list) + metadata: Dict[str, Any] = field(default_factory=dict) + + + @dataclass + class FlowEdge(BaseModel): + """Represents an edge in the control flow graph.""" + + source: str + target: str + edge_type: str = "control" # control, data, call + label: Optional[str] = None + conditions: List[str] = field(default_factory=list) + + + @dataclass + class FunctionInfo(BaseModel): + """Information about a function/method.""" + + name: str + qualified_name: str + file: str + line: int + column: int = 0 + module: str = "" + class_name: Optional[str] = None + is_method: bool = False + is_private: bool = False + is_property: bool = False + docstring: Optional[str] = None + args: List[str] = field(default_factory=list) + returns: Optional[str] = None + decorators: List[str] = field(default_factory=list) + + # CFG info + cfg_entry: Optional[str] = None + cfg_exit: Optional[str] = None + cfg_nodes: List[str] = field(default_factory=list) + calls: List[str] = field(default_factory=list) + called_by: List[str] = field(default_factory=list) + + # Advanced metrics + complexity: Dict[str, Any] = field(default_factory=dict) # Cyclomatic, Cognitive + centrality: float = 0.0 # Betweenness Centrality + reachability: str = "unknown" # reachable, unreachable, unknown + + + @dataclass + class ClassInfo(BaseModel): + """Information about a class.""" + + name: str + qualified_name: str + file: str + line: int + module: str = "" + bases: List[str] = field(default_factory=list) + methods: List[str] = field(default_factory=list) + docstring: Optional[str] = None + is_state_machine: bool = False + + + @dataclass + class ModuleInfo(BaseModel): + """Information about a module/package.""" + + name: str + file: str + is_package: bool = False + source_kind: str = "source" + line_count: int = 0 + imports: List[str] = field(default_factory=list) + functions: List[str] = field(default_factory=list) + classes: List[str] = field(default_factory=list) + + + @dataclass + class Pattern(BaseModel): + """Detected behavioral pattern.""" + + name: str + type: str # recursion, state_machine, factory, singleton, strategy, loop + confidence: float # 0.0 to 1.0 + functions: List[str] = field(default_factory=list) + entry_points: List[str] = field(default_factory=list) + exit_points: List[str] = field(default_factory=list) + metadata: Dict[str, Any] = field(default_factory=dict) + + + @dataclass + class CodeSmell(BaseModel): + """Represents a detected code smell.""" + + name: str + type: str # god_function, feature_envy, etc. + file: str + line: int + severity: float # 0.0 to 1.0 + description: str + context: Dict[str, Any] = field(default_factory=dict) + + + @dataclass + class Mutation(BaseModel): + """Represents a mutation of a variable/object.""" + + variable: str + file: str + line: int + type: str # assign, aug_assign, method_call + scope: str + context: str + + + @dataclass + class DataFlow(BaseModel): + """Represents data flow for a variable.""" + + variable: str + dependencies: Set[str] = field(default_factory=set) + metadata: Dict[str, Any] = field(default_factory=dict) + + + @dataclass + class AnalysisResult(BaseModel): + """Complete analysis result for a project.""" + + project_path: str = "" + analysis_mode: str = "static" + stats: Dict[str, int] = field(default_factory=dict) + + # Graph data + nodes: Dict[str, FlowNode] = field(default_factory=dict) + edges: List[FlowEdge] = field(default_factory=list) + + # Code structure + modules: Dict[str, ModuleInfo] = field(default_factory=dict) + classes: Dict[str, ClassInfo] = field(default_factory=dict) + functions: Dict[str, FunctionInfo] = field(default_factory=dict) + + # Analysis results + patterns: List[Pattern] = field(default_factory=list) + call_graph: Dict[str, List[str]] = field(default_factory=dict) + entry_points: List[str] = field(default_factory=list) + data_flows: Dict[str, DataFlow] = field(default_factory=dict) + + # Refactoring data + metrics: Dict[str, Dict[str, Any]] = field(default_factory=dict) + smells: List[CodeSmell] = field(default_factory=list) + coupling: Dict[str, Any] = field(default_factory=dict) + mutations: List[Mutation] = field(default_factory=list) + + def get_function_count(self) -> int: + """Get total function count.""" + return len(self.functions) + + def get_class_count(self) -> int: + """Get total class count.""" + return len(self.classes) + + def get_node_count(self) -> int: + """Get total CFG node count.""" + return len(self.nodes) + + def get_edge_count(self) -> int: + """Get total edge count.""" + return len(self.edges) diff --git a/docs/refactoring/modularization-and-rust-engine-plan.md b/docs/refactoring/modularization-and-rust-engine-plan.md index 89cb2f1..534bffe 100644 --- a/docs/refactoring/modularization-and-rust-engine-plan.md +++ b/docs/refactoring/modularization-and-rust-engine-plan.md @@ -155,14 +155,16 @@ graph TD - Initial `code2llm-rust` package with PyO3, Rayon, Brandes' centrality (42.3x speedup), CC, and calls. - `native_bridge.py` in `code2llm`. -2. **Phase 2 (Tarjan SCC & Native Reachability in `code2llm-rust`)**: - - Implement Tarjan's Strongly Connected Components in `code2llm-rust` to eliminate NetworkX cycle detection freeze on >1,000 nodes. - - Implement in-memory graph reachability in Rust, completely removing the slow `vulture` file-scanning pass. - -3. **Phase 3 (Native Multi-Language Tree-sitter & File Walking)**: - - Integrate `walkdir` + `ignore` into `code2llm-rust` for instant parallel file collection. - - Embed Tree-sitter grammars in Rust for Python, TypeScript, and Go. - -4. **Phase 4 (Repository Split)**: - - Extract `code2flow` into standalone package to drop `matplotlib` and image rendering from the core analysis path. - - Package `code2toon` as an independent standard for repository state interchange. +2. **Phase 2 (Complete)**: + - Implemented Tarjan's Strongly Connected Components in `code2llm-rust` eliminating NetworkX cycle detection freeze. + - Implemented fast in-memory graph reachability and module coupling ($Ca, Ce, I$) in Rust. + +3. **Phase 3 (Complete)**: + - Integrated `ignore` crate into `code2llm-rust` for instant parallel gitignore-aware file collection (`native_walk_project_files`). + - Implemented parallel cache validation (`PersistentCache`) with Rayon & SHA-256 in Rust. + - Implemented native call graph resolution with $O(1)$ candidate lookups and module affinity matching. + +4. **Phase 4 (In Progress - Ecosystem Modularization)**: + - **`code2graph` (Complete)**: Extracted graph data models (`AnalysisResult`, `FunctionInfo`, `FlowNode`, `FlowEdge`, `GraphQuery`) into standalone package at `packages/code2graph` and integrated into `code2llm/core/models.py`. + - **`code2toon` (Next)**: Extract TOON format specification, serializer, and validator into standalone package. + - **`code2flow`**: Extract visual diagrams and heavy rendering dependencies into standalone package. diff --git a/packages/code2graph/README.md b/packages/code2graph/README.md new file mode 100644 index 0000000..1a0e94c --- /dev/null +++ b/packages/code2graph/README.md @@ -0,0 +1,36 @@ +# code2graph + +Pure domain data models and graph query abstractions for static code analysis, extracted from `code2llm`. + +## Overview + +`code2graph` defines standard, language-agnostic data structures to represent: +- **Control Flow & Data Flow**: `FlowNode`, `FlowEdge`, `DataFlow`, `Mutation` +- **Source Code Structure**: `ModuleInfo`, `ClassInfo`, `FunctionInfo` +- **Analysis & Health**: `AnalysisResult`, `Pattern`, `CodeSmell` +- **Graph Queries**: In-memory query API to inspect call chains, entry points, and reachability. + +## Installation + +```bash +pip install code2graph +``` + +## Usage + +```python +from code2graph import AnalysisResult, FunctionInfo, FlowEdge + +result = AnalysisResult(project_path=".") +func = FunctionInfo( + name="main", + qualified_name="app.main", + file="app.py", + line=1, +) +result.functions["app.main"] = func +``` + +## License + +Apache-2.0 diff --git a/packages/code2graph/code2graph/__init__.py b/packages/code2graph/code2graph/__init__.py new file mode 100644 index 0000000..4e9a51e --- /dev/null +++ b/packages/code2graph/code2graph/__init__.py @@ -0,0 +1,34 @@ +"""code2graph: Domain data models and graph representation for code analysis.""" + +from .models import ( + BaseModel, + FlowNode, + FlowEdge, + FunctionInfo, + ClassInfo, + ModuleInfo, + Pattern, + CodeSmell, + Mutation, + DataFlow, + AnalysisResult, +) +from .query import GraphQuery + +__version__ = "0.1.0" + +__all__ = [ + "BaseModel", + "FlowNode", + "FlowEdge", + "FunctionInfo", + "ClassInfo", + "ModuleInfo", + "Pattern", + "CodeSmell", + "Mutation", + "DataFlow", + "AnalysisResult", + "GraphQuery", + "__version__", +] diff --git a/packages/code2graph/code2graph/models.py b/packages/code2graph/code2graph/models.py new file mode 100644 index 0000000..f5d8f99 --- /dev/null +++ b/packages/code2graph/code2graph/models.py @@ -0,0 +1,207 @@ +from dataclasses import dataclass, field, asdict +from typing import List, Dict, Set, Optional, Any + + +class BaseModel: + """Base class for models with automated serialization.""" + + def to_dict(self, compact: bool = True) -> dict: + """Convert to dictionary using dataclasses.asdict with filtering.""" + data = asdict(self) + if compact: + return self._filter_compact(data) + return data + + def _filter_compact(self, data: Any) -> Any: + """Recursively filter out None and empty collections if compact.""" + if isinstance(data, dict): + return { + k: self._filter_compact(v) + for k, v in data.items() + if v is not None + and (not isinstance(v, (list, dict, set)) or len(v) > 0) + } + elif isinstance(data, (list, tuple, set)): + return [self._filter_compact(v) for v in data] + return data + + +@dataclass +class FlowNode(BaseModel): + """Represents a node in the control flow graph.""" + + id: str + type: str # FUNC, CALL, IF, FOR, WHILE, ASSIGN, RETURN, ENTRY, EXIT + label: str + function: Optional[str] = None + file: Optional[str] = None + line: Optional[int] = None + column: Optional[int] = None + conditions: List[str] = field(default_factory=list) + data_flow: List[str] = field(default_factory=list) + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class FlowEdge(BaseModel): + """Represents an edge in the control flow graph.""" + + source: str + target: str + edge_type: str = "control" # control, data, call + label: Optional[str] = None + conditions: List[str] = field(default_factory=list) + + +@dataclass +class FunctionInfo(BaseModel): + """Information about a function/method.""" + + name: str + qualified_name: str + file: str + line: int + column: int = 0 + module: str = "" + class_name: Optional[str] = None + is_method: bool = False + is_private: bool = False + is_property: bool = False + docstring: Optional[str] = None + args: List[str] = field(default_factory=list) + returns: Optional[str] = None + decorators: List[str] = field(default_factory=list) + + # CFG info + cfg_entry: Optional[str] = None + cfg_exit: Optional[str] = None + cfg_nodes: List[str] = field(default_factory=list) + calls: List[str] = field(default_factory=list) + called_by: List[str] = field(default_factory=list) + + # Advanced metrics + complexity: Dict[str, Any] = field(default_factory=dict) # Cyclomatic, Cognitive + centrality: float = 0.0 # Betweenness Centrality + reachability: str = "unknown" # reachable, unreachable, unknown + + +@dataclass +class ClassInfo(BaseModel): + """Information about a class.""" + + name: str + qualified_name: str + file: str + line: int + module: str = "" + bases: List[str] = field(default_factory=list) + methods: List[str] = field(default_factory=list) + docstring: Optional[str] = None + is_state_machine: bool = False + + +@dataclass +class ModuleInfo(BaseModel): + """Information about a module/package.""" + + name: str + file: str + is_package: bool = False + source_kind: str = "source" + line_count: int = 0 + imports: List[str] = field(default_factory=list) + functions: List[str] = field(default_factory=list) + classes: List[str] = field(default_factory=list) + + +@dataclass +class Pattern(BaseModel): + """Detected behavioral pattern.""" + + name: str + type: str # recursion, state_machine, factory, singleton, strategy, loop + confidence: float # 0.0 to 1.0 + functions: List[str] = field(default_factory=list) + entry_points: List[str] = field(default_factory=list) + exit_points: List[str] = field(default_factory=list) + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class CodeSmell(BaseModel): + """Represents a detected code smell.""" + + name: str + type: str # god_function, feature_envy, etc. + file: str + line: int + severity: float # 0.0 to 1.0 + description: str + context: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class Mutation(BaseModel): + """Represents a mutation of a variable/object.""" + + variable: str + file: str + line: int + type: str # assign, aug_assign, method_call + scope: str + context: str + + +@dataclass +class DataFlow(BaseModel): + """Represents data flow for a variable.""" + + variable: str + dependencies: Set[str] = field(default_factory=set) + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class AnalysisResult(BaseModel): + """Complete analysis result for a project.""" + + project_path: str = "" + analysis_mode: str = "static" + stats: Dict[str, int] = field(default_factory=dict) + + # Graph data + nodes: Dict[str, FlowNode] = field(default_factory=dict) + edges: List[FlowEdge] = field(default_factory=list) + + # Code structure + modules: Dict[str, ModuleInfo] = field(default_factory=dict) + classes: Dict[str, ClassInfo] = field(default_factory=dict) + functions: Dict[str, FunctionInfo] = field(default_factory=dict) + + # Analysis results + patterns: List[Pattern] = field(default_factory=list) + call_graph: Dict[str, List[str]] = field(default_factory=dict) + entry_points: List[str] = field(default_factory=list) + data_flows: Dict[str, DataFlow] = field(default_factory=dict) + + # Refactoring data + metrics: Dict[str, Dict[str, Any]] = field(default_factory=dict) + smells: List[CodeSmell] = field(default_factory=list) + coupling: Dict[str, Any] = field(default_factory=dict) + mutations: List[Mutation] = field(default_factory=list) + + def get_function_count(self) -> int: + """Get total function count.""" + return len(self.functions) + + def get_class_count(self) -> int: + """Get total class count.""" + return len(self.classes) + + def get_node_count(self) -> int: + """Get total CFG node count.""" + return len(self.nodes) + + def get_edge_count(self) -> int: + """Get total edge count.""" + return len(self.edges) diff --git a/packages/code2graph/code2graph/query.py b/packages/code2graph/code2graph/query.py new file mode 100644 index 0000000..3077228 --- /dev/null +++ b/packages/code2graph/code2graph/query.py @@ -0,0 +1,53 @@ +"""Graph query abstractions for code2graph.""" + +from typing import List, Optional, Set +from .models import AnalysisResult, FunctionInfo + + +class GraphQuery: + """Convenience queries over an AnalysisResult.""" + + def __init__(self, result: AnalysisResult): + self.result = result + + def get_function(self, qname: str) -> Optional[FunctionInfo]: + """Look up function by qualified name.""" + return self.result.functions.get(qname) + + def get_callees(self, qname: str) -> List[str]: + """Return list of functions called by qname.""" + func = self.get_function(qname) + return list(func.calls) if func else [] + + def get_callers(self, qname: str) -> List[str]: + """Return list of functions that call qname.""" + func = self.get_function(qname) + return list(func.called_by) if func else [] + + def get_entry_points(self) -> List[str]: + """Return functions with zero callers.""" + if self.result.entry_points: + return list(self.result.entry_points) + return [ + name for name, f in self.result.functions.items() + if not f.called_by + ] + + def find_reachable(self, start_nodes: List[str]) -> Set[str]: + """Find all functions transitively reachable from start_nodes via BFS.""" + visited: Set[str] = set() + queue = list(start_nodes) + for node in start_nodes: + visited.add(node) + + while queue: + curr = queue.pop(0) + func = self.get_function(curr) + if not func: + continue + for callee in func.calls: + if callee not in visited and callee in self.result.functions: + visited.add(callee) + queue.append(callee) + + return visited diff --git a/packages/code2graph/pyproject.toml b/packages/code2graph/pyproject.toml new file mode 100644 index 0000000..09da7c5 --- /dev/null +++ b/packages/code2graph/pyproject.toml @@ -0,0 +1,26 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "code2graph" +version = "0.1.0" +description = "Semantic code graph and topology models for static analysis" +readme = "README.md" +requires-python = ">=3.9" +license = {text = "Apache-2.0"} +authors = [ + {name = "Tom Sapletta", email = "tom@sapletta.com"} +] +keywords = ["code-analysis", "graph", "ast", "call-graph", "static-analysis"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3", + "Topic :: Software Development :: Libraries :: Python Modules", +] +dependencies = [] + +[tool.setuptools.packages.find] +where = ["."] +include = ["code2graph*"] diff --git a/packages/code2graph/tests/test_code2graph.py b/packages/code2graph/tests/test_code2graph.py new file mode 100644 index 0000000..be6416f --- /dev/null +++ b/packages/code2graph/tests/test_code2graph.py @@ -0,0 +1,42 @@ +"""Tests for code2graph data models and queries.""" + +from code2graph import ( + AnalysisResult, + FunctionInfo, + ClassInfo, + ModuleInfo, + FlowNode, + FlowEdge, + GraphQuery, +) + + +def test_models_serialization(): + result = AnalysisResult(project_path="/tmp/test") + fn = FunctionInfo( + name="test_func", + qualified_name="pkg.mod.test_func", + file="pkg/mod.py", + line=10, + calls=["pkg.mod.helper"], + ) + result.functions[fn.qualified_name] = fn + + d = result.to_dict(compact=True) + assert d["project_path"] == "/tmp/test" + assert "pkg.mod.test_func" in d["functions"] + assert d["functions"]["pkg.mod.test_func"]["name"] == "test_func" + + +def test_graph_query(): + result = AnalysisResult() + f1 = FunctionInfo(name="main", qualified_name="app.main", file="app.py", line=1, calls=["app.worker"]) + f2 = FunctionInfo(name="worker", qualified_name="app.worker", file="app.py", line=10, called_by=["app.main"]) + result.functions["app.main"] = f1 + result.functions["app.worker"] = f2 + + query = GraphQuery(result) + assert query.get_callees("app.main") == ["app.worker"] + assert query.get_callers("app.worker") == ["app.main"] + assert query.get_entry_points() == ["app.main"] + assert query.find_reachable(["app.main"]) == {"app.main", "app.worker"} diff --git a/packages/code2llm-rust/src/calls.rs b/packages/code2llm-rust/src/calls.rs index a5821d1..64873cb 100644 --- a/packages/code2llm-rust/src/calls.rs +++ b/packages/code2llm-rust/src/calls.rs @@ -141,18 +141,21 @@ pub fn resolve_call_graph( .entry(resolved.to_string()) .or_default() .push(func_name.clone()); + } else { + res_calls.push(called.clone()); } } resolved_calls.insert(func_name.clone(), res_calls); } let mut entry_points: Vec = Vec::new(); - for (name, callers) in &called_by { - if callers.is_empty() { - entry_points.push(name.clone()); + for (name, _) in functions_with_calls { + if let Some(callers) = called_by.get(name) { + if callers.is_empty() { + entry_points.push(name.clone()); + } } } - entry_points.sort(); CallGraphResolution { resolved_calls, diff --git a/pyproject.toml b/pyproject.toml index a020822..d0ccec3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,6 +68,9 @@ license = "Apache-2.0" native = [ "code2llm-rust>=0.1.0", ] +graph = [ + "code2graph>=0.1.0", +] dev = [ "pytest>=6.2", "pytest-cov>=2.12", @@ -77,6 +80,7 @@ dev = [ "mypy>=0.910", "costs>=0.1.20", "code2llm-rust>=0.1.0", + "code2graph>=0.1.0", ] [project.scripts]