"
+
+ def __hash__(self):
+ return hash((self.start_byte, self.end_byte, self.file_path, self.priority, self.new_content))
+
+ def __eq__(self, other):
+ if not isinstance(other, type(self)):
+ return False
+
+ # Check for everything EXCEPT transaction_id
+ return (
+ self.start_byte == other.start_byte
+ and self.end_byte == other.end_byte
+ and self.file_path == other.file_path
+ and self.priority == other.priority
+ and self._new_content == other._new_content
+ )
+
+ @property
+ def length(self):
+ """Length of the transaction in bytes."""
+ return self.end_byte - self.start_byte
+
+ def execute(self):
+ """Execute the transaction to modify the file."""
+ msg = "Transaction.execute() must be implemented by subclasses"
+ raise NotImplementedError(msg)
+
+ def get_diff(self) -> DiffLite:
+ """Gets the diff produced by this transaction."""
+ msg = "Transaction.get_diff() must be implemented by subclasses"
+ raise NotImplementedError(msg)
+
+ def diff_str(self):
+ """Human-readable string representation of the change."""
+ msg = "Transaction.diff_str() must be implemented by subclasses"
+ raise NotImplementedError(msg)
+
+ def _to_sort_key(transaction: "Transaction"):
+ """Key function for sorting transactions."""
+ # Sort by:
+ # 1. Descending start_byte
+ # 2. Ascending transaction type
+ # 3. Ascending priority
+ # 4. Descending time of transaction
+ priority = (transaction.priority,) if isinstance(transaction.priority, int) else transaction.priority
+
+ return -transaction.start_byte, transaction.transaction_order.value, priority, -transaction.transaction_id
+
+ @cached_property
+ def new_content(self) -> Optional[str]:
+ """Get the new content, evaluating the content function if necessary."""
+ return self._new_content() if isinstance(self._new_content, ContentFunc) else self._new_content
+
+ @staticmethod
+ def create_new_file(filepath: Union[str, Path], content: str) -> "FileAddTransaction":
+ """Create a transaction to add a new file."""
+ return FileAddTransaction(Path(filepath))
+
+ @staticmethod
+ def delete_file(filepath: Union[str, Path]) -> "FileRemoveTransaction":
+ """Create a transaction to delete a file."""
+ # In a real implementation, this would need a File object
+ # For now, we'll create a placeholder implementation
+ from pathlib import Path
+ class FilePlaceholder:
+ def __init__(self, path):
+ self.path = Path(path)
+
+ return FileRemoveTransaction(FilePlaceholder(filepath))
+
+class RemoveTransaction(Transaction):
+ """Transaction to remove content from a file."""
+ transaction_order = TransactionPriority.Remove
+
+ exec_func: Optional[Callable[[], None]] = None
+
+ def __init__(self, start_byte: int, end_byte: int, file: Any, priority: int = 0, exec_func: Optional[Callable[[], None]] = None) -> None:
+ super().__init__(start_byte, end_byte, file.path, priority=priority)
+ self.file = file
+ self.exec_func = exec_func
+
+ def _generate_new_content_bytes(self) -> bytes:
+ """Generate the new content bytes after removal."""
+ content_bytes = self.file.content_bytes
+ new_content_bytes = content_bytes[: self.start_byte] + content_bytes[self.end_byte :]
+ return new_content_bytes
+
+ def execute(self) -> None:
+ """Removes the content between start_byte and end_byte."""
+ self.file.write_bytes(self._generate_new_content_bytes())
+ if self.exec_func:
+ self.exec_func()
+
+ def get_diff(self) -> DiffLite:
+ """Gets the diff produced by this transaction."""
+ return DiffLite(ChangeType.Modified, self.file_path, old_content=self.file.content_bytes)
+
+ def diff_str(self) -> str:
+ """Human-readable string representation of the change."""
+ diff = "".join(unified_diff(self.file.content.splitlines(True), self._generate_new_content_bytes().decode("utf-8").splitlines(True)))
+ return f"Remove {self.length} bytes at bytes ({self.start_byte}, {self.end_byte})\n{diff}"
+
+class InsertTransaction(Transaction):
+ """Transaction to insert content into a file."""
+ transaction_order = TransactionPriority.Insert
+
+ exec_func: Optional[Callable[[], None]] = None
+
+ def __init__(
+ self,
+ insert_byte: int,
+ file: Any,
+ new_content: Union[str, Callable[[], str]],
+ *,
+ priority: Union[int, tuple] = 0,
+ exec_func: Optional[Callable[[], None]] = None,
+ ) -> None:
+ super().__init__(insert_byte, insert_byte, file.path, priority=priority, new_content=new_content)
+ self.insert_byte = insert_byte
+ self.file = file
+ self.exec_func = exec_func
+
+ def _generate_new_content_bytes(self) -> bytes:
+ """Generate the new content bytes after insertion."""
+ if self.new_content is None:
+ raise ValueError("Cannot generate content bytes: new_content is None")
+ new_bytes = bytes(self.new_content, encoding="utf-8")
+ content_bytes = self.file.content_bytes
+ head = content_bytes[: self.insert_byte]
+ tail = content_bytes[self.insert_byte :]
+ new_content_bytes = head + new_bytes + tail
+ return new_content_bytes
+
+ def execute(self) -> None:
+ """Inserts new_src at the specified byte_index."""
+ self.file.write_bytes(self._generate_new_content_bytes())
+ if self.exec_func:
+ self.exec_func()
+
+ def get_diff(self) -> DiffLite:
+ """Gets the diff produced by this transaction."""
+ return DiffLite(ChangeType.Modified, self.file_path, old_content=self.file.content_bytes)
+
+ def diff_str(self) -> str:
+ """Human-readable string representation of the change."""
+ diff = "".join(unified_diff(self.file.content.splitlines(True), self._generate_new_content_bytes().decode("utf-8").splitlines(True)))
+ content_length = len(self.new_content) if self.new_content is not None else 0
+ return f"Insert {content_length} bytes at bytes ({self.start_byte}, {self.end_byte})\n{diff}"
+
+class EditTransaction(Transaction):
+ """Transaction to edit content in a file."""
+ transaction_order = TransactionPriority.Edit
+ new_content: str
+
+ def __init__(
+ self,
+ start_byte: int,
+ end_byte: int,
+ file: Any,
+ new_content: str,
+ priority: int = 0,
+ ) -> None:
+ super().__init__(start_byte, end_byte, file.path, priority=priority, new_content=new_content)
+ self.file = file
+
+ def _generate_new_content_bytes(self) -> bytes:
+ """Generate the new content bytes after editing."""
+ new_bytes = bytes(self.new_content, "utf-8")
+ content_bytes = self.file.content_bytes
+ new_content_bytes = content_bytes[: self.start_byte] + new_bytes + content_bytes[self.end_byte :]
+ return new_content_bytes
+
+ def execute(self) -> None:
+ """Edits the entirety of this node's source to new_src."""
+ self.file.write_bytes(self._generate_new_content_bytes())
+
+ def get_diff(self) -> DiffLite:
+ """Gets the diff produced by this transaction."""
+ return DiffLite(ChangeType.Modified, self.file_path, old_content=self.file.content_bytes)
+
+ def diff_str(self) -> str:
+ """Human-readable string representation of the change."""
+ diff = "".join(unified_diff(self.file.content.splitlines(True), self._generate_new_content_bytes().decode("utf-8").splitlines(True)))
+ return f"Edit {self.length} bytes at bytes ({self.start_byte}, {self.end_byte}), src: ({self.new_content[:50]})\n{diff}"
+
+ def break_down(self) -> Optional[list[InsertTransaction]]:
+ """Break down an edit transaction into insert transactions."""
+ old = self.file.content_bytes[self.start_byte : self.end_byte]
+ new = bytes(self.new_content, "utf-8")
+ if old and old in new:
+ prefix, suffix = new.split(old, maxsplit=1)
+ ret = []
+ if suffix:
+ ret.append(InsertTransaction(self.end_byte, self.file, suffix.decode("utf-8"), priority=self.priority))
+ if prefix:
+ ret.append(InsertTransaction(self.start_byte, self.file, prefix.decode("utf-8"), priority=self.priority))
+ return ret
+ return None
+
+class FileAddTransaction(Transaction):
+ """Transaction to add a new file."""
+ transaction_order = TransactionPriority.FileAdd
+
+ def __init__(
+ self,
+ file_path: Path,
+ priority: int = 0,
+ ) -> None:
+ super().__init__(0, 0, file_path, priority=priority)
+
+ def execute(self) -> None:
+ """Adds a new file."""
+ pass # execute is a no-op as the file is immediately added
+
+ def get_diff(self) -> DiffLite:
+ """Gets the diff produced by this transaction."""
+ return DiffLite(ChangeType.Added, self.file_path)
+
+ def diff_str(self) -> str:
+ """Human-readable string representation of the change."""
+ return f"Add file at {self.file_path}"
+
+class FileRenameTransaction(Transaction):
+ """Transaction to rename a file."""
+ transaction_order = TransactionPriority.FileRename
+
+ def __init__(
+ self,
+ file: Any,
+ new_file_path: str,
+ priority: int = 0,
+ ) -> None:
+ super().__init__(0, 0, file.path, priority=priority, new_content=new_file_path)
+ self.new_file_path = file.ctx.to_absolute(new_file_path) if hasattr(file, 'ctx') else Path(new_file_path)
+ self.file = file
+
+ def execute(self) -> None:
+ """Renames the file."""
+ if hasattr(self.file, 'ctx') and hasattr(self.file.ctx, 'io'):
+ self.file.ctx.io.save_files({self.file.path})
+ self.file_path.rename(self.new_file_path)
+
+ def get_diff(self) -> DiffLite:
+ """Gets the diff produced by this transaction."""
+ return DiffLite(ChangeType.Renamed, self.file_path, self.file_path, self.new_file_path,
+ old_content=self.file.content_bytes if hasattr(self.file, 'content_bytes') else None)
+
+ def diff_str(self) -> str:
+ """Human-readable string representation of the change."""
+ return f"Rename file from {self.file_path} to {self.new_file_path}"
+
+class FileRemoveTransaction(Transaction):
+ """Transaction to remove a file."""
+ transaction_order = TransactionPriority.FileRemove
+
+ def __init__(
+ self,
+ file: Any,
+ priority: int = 0,
+ ) -> None:
+ super().__init__(0, 0, file.path, priority=priority)
+ self.file = file
+
+ def execute(self) -> None:
+ """Removes the file."""
+ if hasattr(self.file, 'ctx') and hasattr(self.file.ctx, 'io'):
+ self.file.ctx.io.delete_file(self.file.path)
+ else:
+ # Fallback for when ctx.io is not available
+ import os
+ if os.path.exists(self.file_path):
+ os.remove(self.file_path)
+
+ def get_diff(self) -> DiffLite:
+ """Gets the diff produced by this transaction."""
+ return DiffLite(ChangeType.Removed, self.file_path,
+ old_content=self.file.content_bytes if hasattr(self.file, 'content_bytes') else None)
+
+ def diff_str(self) -> str:
+ """Human-readable string representation of the change."""
+ return f"Remove file at {self.file_path}"
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/utils.py b/codegen-on-oss/codegen_on_oss/analyzers/utils.py
new file mode 100644
index 000000000..b04da2f70
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/utils.py
@@ -0,0 +1,80 @@
+"""Utility functions for documentation generation."""
+
+import re
+import textwrap
+
+
+def sanitize_docstring_for_markdown(docstring: str | None) -> str:
+ """Sanitize the docstring for MDX.
+
+ Args:
+ docstring: The docstring to sanitize.
+
+ Returns:
+ The sanitized docstring.
+ """
+ if docstring is None:
+ return ""
+ docstring_lines = docstring.splitlines()
+ if len(docstring_lines) > 1:
+ docstring_lines[1:] = [textwrap.dedent(line) for line in docstring_lines[1:]]
+ docstring = "\n".join(docstring_lines)
+ if docstring.startswith('"""'):
+ docstring = docstring[3:]
+ if docstring.endswith('"""'):
+ docstring = docstring[:-3]
+ return docstring
+
+
+def sanitize_mdx_mintlify_description(content: str) -> str:
+ """Mintlify description field needs to have string escaped, which content doesn't need.
+
+ Args:
+ content: The content to sanitize.
+
+ Returns:
+ The sanitized content.
+ """
+ content = sanitize_docstring_for_markdown(content)
+ # make sure all `< />` components are properly escaped with a `` inline-block
+ # if the string already has the single-quote then this is a no-op
+ content = re.sub(r"(?]+>)(?!`)", r"`\1`", content)
+
+ # escape double quote characters
+ if re.search(r'\\"', content):
+ return content # No-op if already escaped
+ return re.sub(r'(")', r"\\\1", content)
+
+
+def sanitize_html_for_mdx(html_string: str) -> str:
+ """Sanitize HTML string for MDX by escaping double quotes in attribute values.
+
+ Args:
+ html_string: The input HTML string to sanitize
+
+ Returns:
+ The sanitized HTML string with escaped quotes
+ """
+ # Replace double quotes with " but only in HTML attributes
+ return re.sub(r'"', """, html_string)
+
+
+def extract_class_description(docstring: str) -> str:
+ """Extract the class description from a docstring, excluding the attributes section.
+
+ Args:
+ docstring: The class docstring to parse
+
+ Returns:
+ The class description with whitespace normalized
+ """
+ if not docstring:
+ return ""
+
+ # Split by "Attributes:" and take only the first part
+ parts = docstring.split("Attributes:")
+ description = parts[0]
+
+ # Normalize whitespace
+ lines = [line.strip() for line in description.strip().splitlines()]
+ return " ".join(filter(None, lines))
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/__init__.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/analysis_visualizer.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/analysis_visualizer.py
new file mode 100644
index 000000000..a7c2a3f77
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/analysis_visualizer.py
@@ -0,0 +1,558 @@
+#!/usr/bin/env python3
+"""
+Analysis Visualizer Module
+
+This module provides visualization capabilities for code analysis results
+including dead code detection, cyclomatic complexity, and issue heatmaps.
+"""
+
+import logging
+
+from .visualizer import BaseVisualizer, OutputFormat, VisualizationType
+
+try:
+ import matplotlib.pyplot as plt
+ import networkx as nx
+ from matplotlib.colors import LinearSegmentedColormap
+except ImportError:
+ logging.warning(
+ "Visualization dependencies not found. Please install them with: pip install networkx matplotlib"
+ )
+
+logger = logging.getLogger(__name__)
+
+
+class AnalysisVisualizer(BaseVisualizer):
+ """
+ Visualizer for code analysis results.
+
+ This class provides methods to visualize analysis results such as
+ dead code detection, cyclomatic complexity, and issue heatmaps.
+ """
+
+ def __init__(self, analyzer=None, codebase=None, context=None, **kwargs):
+ """
+ Initialize the AnalysisVisualizer.
+
+ Args:
+ analyzer: Analyzer with analysis results
+ codebase: Codebase instance to visualize
+ context: Context providing graph representation
+ **kwargs: Additional configuration options
+ """
+ super().__init__(**kwargs)
+ self.analyzer = analyzer
+ self.codebase = codebase or (analyzer.base_codebase if analyzer else None)
+ self.context = context or (analyzer.base_context if analyzer else None)
+
+ def visualize_dead_code(self, path_filter: str | None = None):
+ """
+ Generate a visualization of dead (unused) code in the codebase.
+
+ Args:
+ path_filter: Optional path to filter files
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ entity_name = path_filter or "codebase"
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Check for analyzer
+ if not self.analyzer:
+ logger.error("Analyzer required for dead code visualization")
+ return None
+
+ # Check for analysis results
+ if not hasattr(self.analyzer, "results") or not self.analyzer.results:
+ logger.error("Analysis results not available")
+ return None
+
+ # Extract dead code information from analysis results
+ dead_code = {}
+ if (
+ "static_analysis" in self.analyzer.results
+ and "dead_code" in self.analyzer.results["static_analysis"]
+ ):
+ dead_code = self.analyzer.results["static_analysis"]["dead_code"]
+
+ if not dead_code:
+ logger.warning("No dead code detected in analysis results")
+ return None
+
+ # Create file nodes for containing dead code
+ file_nodes = {}
+
+ # Process unused functions
+ if "unused_functions" in dead_code:
+ for unused_func in dead_code["unused_functions"]:
+ file_path = unused_func.get("file", "")
+
+ # Skip if path filter is specified and doesn't match
+ if path_filter and not file_path.startswith(path_filter):
+ continue
+
+ # Add file node if not already added
+ if file_path not in file_nodes:
+ # Find file in codebase
+ file_obj = None
+ for file in self.codebase.files:
+ if hasattr(file, "path") and str(file.path) == file_path:
+ file_obj = file
+ break
+
+ if file_obj:
+ file_name = file_path.split("/")[-1]
+ self._add_node(
+ file_obj,
+ name=file_name,
+ color=self.config.color_palette.get("File"),
+ file_path=file_path,
+ )
+
+ file_nodes[file_path] = file_obj
+
+ # Add unused function node
+ func_name = unused_func.get("name", "")
+ func_line = unused_func.get("line", None)
+
+ # Create a placeholder for the function (we don't have the actual object)
+ func_obj = {
+ "name": func_name,
+ "file_path": file_path,
+ "line": func_line,
+ "type": "Function",
+ }
+
+ self._add_node(
+ func_obj,
+ name=func_name,
+ color=self.config.color_palette.get("Dead"),
+ file_path=file_path,
+ line=func_line,
+ is_dead=True,
+ )
+
+ # Add edge from file to function
+ if file_path in file_nodes:
+ self._add_edge(
+ file_nodes[file_path], func_obj, type="contains_dead"
+ )
+
+ # Process unused variables
+ if "unused_variables" in dead_code:
+ for unused_var in dead_code["unused_variables"]:
+ file_path = unused_var.get("file", "")
+
+ # Skip if path filter is specified and doesn't match
+ if path_filter and not file_path.startswith(path_filter):
+ continue
+
+ # Add file node if not already added
+ if file_path not in file_nodes:
+ # Find file in codebase
+ file_obj = None
+ for file in self.codebase.files:
+ if hasattr(file, "path") and str(file.path) == file_path:
+ file_obj = file
+ break
+
+ if file_obj:
+ file_name = file_path.split("/")[-1]
+ self._add_node(
+ file_obj,
+ name=file_name,
+ color=self.config.color_palette.get("File"),
+ file_path=file_path,
+ )
+
+ file_nodes[file_path] = file_obj
+
+ # Add unused variable node
+ var_name = unused_var.get("name", "")
+ var_line = unused_var.get("line", None)
+
+ # Create a placeholder for the variable
+ var_obj = {
+ "name": var_name,
+ "file_path": file_path,
+ "line": var_line,
+ "type": "Variable",
+ }
+
+ self._add_node(
+ var_obj,
+ name=var_name,
+ color=self.config.color_palette.get("Dead"),
+ file_path=file_path,
+ line=var_line,
+ is_dead=True,
+ )
+
+ # Add edge from file to variable
+ if file_path in file_nodes:
+ self._add_edge(file_nodes[file_path], var_obj, type="contains_dead")
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.DEAD_CODE, entity_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.DEAD_CODE, entity_name, fig
+ )
+
+ def visualize_cyclomatic_complexity(self, path_filter: str | None = None):
+ """
+ Generate a heatmap visualization of cyclomatic complexity.
+
+ Args:
+ path_filter: Optional path to filter files
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ entity_name = path_filter or "codebase"
+
+ # Check for analyzer
+ if not self.analyzer:
+ logger.error("Analyzer required for complexity visualization")
+ return None
+
+ # Check for analysis results
+ if not hasattr(self.analyzer, "results") or not self.analyzer.results:
+ logger.error("Analysis results not available")
+ return None
+
+ # Extract complexity information from analysis results
+ complexity_data = {}
+ if (
+ "static_analysis" in self.analyzer.results
+ and "code_complexity" in self.analyzer.results["static_analysis"]
+ ):
+ complexity_data = self.analyzer.results["static_analysis"][
+ "code_complexity"
+ ]
+
+ if not complexity_data:
+ logger.warning("No complexity data found in analysis results")
+ return None
+
+ # Extract function complexities
+ functions = []
+ if "function_complexity" in complexity_data:
+ for func_data in complexity_data["function_complexity"]:
+ # Skip if path filter is specified and doesn't match
+ if path_filter and not func_data.get("file", "").startswith(
+ path_filter
+ ):
+ continue
+
+ functions.append({
+ "name": func_data.get("name", ""),
+ "file": func_data.get("file", ""),
+ "complexity": func_data.get("complexity", 1),
+ "line": func_data.get("line", None),
+ })
+
+ # Sort functions by complexity (descending)
+ functions.sort(key=lambda x: x.get("complexity", 0), reverse=True)
+
+ # Generate heatmap visualization
+ plt.figure(figsize=(12, 10))
+
+ # Extract data for heatmap
+ func_names = [
+ f"{func['name']} ({func['file'].split('/')[-1]})" for func in functions[:30]
+ ]
+ complexities = [func.get("complexity", 0) for func in functions[:30]]
+
+ # Create horizontal bar chart
+ bars = plt.barh(func_names, complexities)
+
+ # Color bars by complexity
+ norm = plt.Normalize(1, max(10, max(complexities)))
+ cmap = plt.cm.get_cmap("YlOrRd")
+
+ for i, bar in enumerate(bars):
+ complexity = complexities[i]
+ bar.set_color(cmap(norm(complexity)))
+
+ # Add labels and title
+ plt.xlabel("Cyclomatic Complexity")
+ plt.title("Top Functions by Cyclomatic Complexity")
+ plt.grid(axis="x", linestyle="--", alpha=0.6)
+
+ # Add colorbar
+ plt.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=cmap), label="Complexity")
+
+ # Save and return visualization
+ return self._save_visualization(
+ VisualizationType.CYCLOMATIC_COMPLEXITY, entity_name, plt.gcf()
+ )
+
+ def visualize_issues_heatmap(self, severity=None, path_filter: str | None = None):
+ """
+ Generate a heatmap visualization of issues in the codebase.
+
+ Args:
+ severity: Optional severity level to filter issues
+ path_filter: Optional path to filter files
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ entity_name = f"{severity.value if severity else 'all'}_issues"
+
+ # Check for analyzer
+ if not self.analyzer:
+ logger.error("Analyzer required for issues visualization")
+ return None
+
+ # Check for analysis results
+ if (
+ not hasattr(self.analyzer, "results")
+ or "issues" not in self.analyzer.results
+ ):
+ logger.error("Issues not available in analysis results")
+ return None
+
+ issues = self.analyzer.results["issues"]
+
+ # Filter issues by severity if specified
+ if severity:
+ issues = [issue for issue in issues if issue.get("severity") == severity]
+
+ # Filter issues by path if specified
+ if path_filter:
+ issues = [
+ issue
+ for issue in issues
+ if issue.get("file", "").startswith(path_filter)
+ ]
+
+ if not issues:
+ logger.warning("No issues found matching the criteria")
+ return None
+
+ # Group issues by file
+ file_issues = {}
+ for issue in issues:
+ file_path = issue.get("file", "")
+ if file_path not in file_issues:
+ file_issues[file_path] = []
+
+ file_issues[file_path].append(issue)
+
+ # Generate heatmap visualization
+ plt.figure(figsize=(12, 10))
+
+ # Extract data for heatmap
+ files = list(file_issues.keys())
+ file_names = [file_path.split("/")[-1] for file_path in files]
+ issue_counts = [len(file_issues[file_path]) for file_path in files]
+
+ # Sort by issue count
+ sorted_data = sorted(
+ zip(file_names, issue_counts, files, strict=False),
+ key=lambda x: x[1],
+ reverse=True,
+ )
+ file_names, issue_counts, files = zip(*sorted_data, strict=False)
+
+ # Create horizontal bar chart
+ bars = plt.barh(file_names[:20], issue_counts[:20])
+
+ # Color bars by issue count
+ norm = plt.Normalize(1, max(5, max(issue_counts[:20])))
+ cmap = plt.cm.get_cmap("OrRd")
+
+ for i, bar in enumerate(bars):
+ count = issue_counts[i]
+ bar.set_color(cmap(norm(count)))
+
+ # Add labels and title
+ plt.xlabel("Number of Issues")
+ severity_text = f" ({severity.value})" if severity else ""
+ plt.title(f"Files with the Most Issues{severity_text}")
+ plt.grid(axis="x", linestyle="--", alpha=0.6)
+
+ # Add colorbar
+ plt.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=cmap), label="Issue Count")
+
+ # Save and return visualization
+ return self._save_visualization(
+ VisualizationType.ISSUES_HEATMAP, entity_name, plt.gcf()
+ )
+
+ def visualize_pr_comparison(self):
+ """
+ Generate a visualization comparing base branch with PR.
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ # Check for analyzer with PR data
+ if (
+ not self.analyzer
+ or not hasattr(self.analyzer, "pr_codebase")
+ or not self.analyzer.pr_codebase
+ or not self.analyzer.base_codebase
+ ):
+ logger.error("PR comparison requires analyzer with PR data")
+ return None
+
+ entity_name = (
+ f"pr_{self.analyzer.pr_number}"
+ if hasattr(self.analyzer, "pr_number") and self.analyzer.pr_number
+ else "pr_comparison"
+ )
+
+ # Check for analysis results
+ if (
+ not hasattr(self.analyzer, "results")
+ or "comparison" not in self.analyzer.results
+ ):
+ logger.error("Comparison data not available in analysis results")
+ return None
+
+ comparison = self.analyzer.results["comparison"]
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Process symbol comparison data
+ if "symbol_comparison" in comparison:
+ for symbol_data in comparison["symbol_comparison"]:
+ symbol_name = symbol_data.get("name", "")
+ in_base = symbol_data.get("in_base", False)
+ in_pr = symbol_data.get("in_pr", False)
+
+ # Create a placeholder for the symbol
+ symbol_obj = {
+ "name": symbol_name,
+ "in_base": in_base,
+ "in_pr": in_pr,
+ "type": "Symbol",
+ }
+
+ # Determine node color based on presence in base and PR
+ if in_base and in_pr:
+ color = "#A5D6A7" # Light green (modified)
+ elif in_base:
+ color = "#EF9A9A" # Light red (removed)
+ else:
+ color = "#90CAF9" # Light blue (added)
+
+ # Add node for symbol
+ self._add_node(
+ symbol_obj,
+ name=symbol_name,
+ color=color,
+ in_base=in_base,
+ in_pr=in_pr,
+ )
+
+ # Process parameter changes if available
+ if "parameter_changes" in symbol_data:
+ param_changes = symbol_data["parameter_changes"]
+
+ # Process removed parameters
+ for param in param_changes.get("removed", []):
+ param_obj = {
+ "name": param,
+ "change_type": "removed",
+ "type": "Parameter",
+ }
+
+ self._add_node(
+ param_obj,
+ name=param,
+ color="#EF9A9A", # Light red (removed)
+ change_type="removed",
+ )
+
+ self._add_edge(symbol_obj, param_obj, type="removed_parameter")
+
+ # Process added parameters
+ for param in param_changes.get("added", []):
+ param_obj = {
+ "name": param,
+ "change_type": "added",
+ "type": "Parameter",
+ }
+
+ self._add_node(
+ param_obj,
+ name=param,
+ color="#90CAF9", # Light blue (added)
+ change_type="added",
+ )
+
+ self._add_edge(symbol_obj, param_obj, type="added_parameter")
+
+ # Process return type changes if available
+ if "return_type_change" in symbol_data:
+ return_type_change = symbol_data["return_type_change"]
+ old_type = return_type_change.get("old", "None")
+ new_type = return_type_change.get("new", "None")
+
+ return_obj = {
+ "name": f"{old_type} -> {new_type}",
+ "old_type": old_type,
+ "new_type": new_type,
+ "type": "ReturnType",
+ }
+
+ self._add_node(
+ return_obj,
+ name=f"{old_type} -> {new_type}",
+ color="#FFD54F", # Amber (changed)
+ old_type=old_type,
+ new_type=new_type,
+ )
+
+ self._add_edge(symbol_obj, return_obj, type="return_type_change")
+
+ # Process call site issues if available
+ if "call_site_issues" in symbol_data:
+ for issue in symbol_data["call_site_issues"]:
+ issue_file = issue.get("file", "")
+ issue_line = issue.get("line", None)
+ issue_text = issue.get("issue", "")
+
+ # Create a placeholder for the issue
+ issue_obj = {
+ "name": issue_text,
+ "file": issue_file,
+ "line": issue_line,
+ "type": "Issue",
+ }
+
+ self._add_node(
+ issue_obj,
+ name=f"{issue_file.split('/')[-1]}:{issue_line}",
+ color="#EF5350", # Red (error)
+ file_path=issue_file,
+ line=issue_line,
+ issue_text=issue_text,
+ )
+
+ self._add_edge(symbol_obj, issue_obj, type="call_site_issue")
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.PR_COMPARISON, entity_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.PR_COMPARISON, entity_name, fig
+ )
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/__init__.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/__init__.py
new file mode 100644
index 000000000..e9e9da182
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/__init__.py
@@ -0,0 +1,6 @@
+"""
+Call Graph Visualization Module
+
+This module provides tools for visualizing call graphs and function relationships in a codebase.
+"""
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/call_trace.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/call_trace.py
new file mode 100644
index 000000000..85448ac4f
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/call_trace.py
@@ -0,0 +1,83 @@
+import codegen
+import networkx as nx
+from codegen import Codebase
+from codegen.sdk.core.class_definition import Class
+from codegen.sdk.core.import_resolution import Import
+from codegen.sdk.core.symbol import Symbol
+
+G = nx.DiGraph()
+
+IGNORE_EXTERNAL_MODULE_CALLS = True
+IGNORE_CLASS_CALLS = False
+MAX_DEPTH = 10
+
+COLOR_PALETTE = {
+ "StartFunction": "#9cdcfe", # Light blue for the starting function
+ "PyFunction": "#a277ff", # Purple for Python functions
+ "PyClass": "#ffca85", # Orange for Python classes
+ "ExternalModule": "#f694ff", # Pink for external module references
+}
+
+# Dictionary to track visited nodes and prevent cycles
+visited = {}
+
+
+def create_dependencies_visualization(symbol: Symbol, depth: int = 0):
+ """Creates a visualization of symbol dependencies in the codebase
+
+ Recursively traverses the dependency tree of a symbol (function, class, etc.)
+ and creates a directed graph representation. Dependencies can be either direct
+ symbol references or imports.
+
+ Args:
+ symbol (Symbol): The starting symbol whose dependencies will be mapped
+ depth (int): Current depth in the recursive traversal
+ """
+ if depth >= MAX_DEPTH:
+ return
+
+ for dep in symbol.dependencies:
+ dep_symbol = None
+
+ if isinstance(dep, Symbol):
+ dep_symbol = dep
+ elif isinstance(dep, Import):
+ dep_symbol = dep.resolved_symbol if dep.resolved_symbol else None
+
+ if dep_symbol:
+ G.add_node(dep_symbol, color=COLOR_PALETTE.get(dep_symbol.__class__.__name__, "#f694ff"))
+ G.add_edge(symbol, dep_symbol)
+
+ if not isinstance(dep_symbol, Class):
+ create_dependencies_visualization(dep_symbol, depth + 1)
+
+
+@codegen.function("visualize-symbol-dependencies")
+def run(codebase: Codebase):
+ """Generate a visualization of symbol dependencies in a codebase.
+
+ This codemod:
+ 1. Creates a directed graph of symbol dependencies starting from a target function
+ 2. Tracks relationships between functions, classes, and imports
+ 3. Generates a visual representation of the dependency hierarchy
+ """
+ global G
+ G = nx.DiGraph()
+
+ target_func = codebase.get_function("get_query_runner")
+ G.add_node(target_func, color=COLOR_PALETTE.get("StartFunction"))
+
+ create_dependencies_visualization(target_func)
+
+ print(G)
+ print("Use codegen.sh to visualize the graph!")
+
+
+if __name__ == "__main__":
+ print("Initializing codebase...")
+ codebase = Codebase.from_repo("codegen-oss/posthog", commit="b174f2221ea4ae50e715eb6a7e70e9a2b0760800", language="python")
+ print(f"Codebase with {len(codebase.files)} files and {len(codebase.functions)} functions.")
+ print("Creating graph...")
+
+ run(codebase)
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/graph_viz_call_graph.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/graph_viz_call_graph.py
new file mode 100644
index 000000000..9fd770841
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/graph_viz_call_graph.py
@@ -0,0 +1,358 @@
+from abc import ABC
+
+import networkx as nx
+
+from codegen.sdk.core.class_definition import Class
+from codegen.sdk.core.codebase import CodebaseType
+from codegen.sdk.core.detached_symbols.function_call import FunctionCall
+from codegen.sdk.core.external_module import ExternalModule
+from codegen.sdk.core.function import Function
+from codegen.sdk.core.interfaces.callable import Callable
+from codegen.shared.enums.programming_language import ProgrammingLanguage
+from tests.shared.skills.decorators import skill, skill_impl
+from tests.shared.skills.skill import Skill
+from tests.shared.skills.skill_test import SkillTestCase, SkillTestCasePyFile
+
+CallGraphFromNodeTest = SkillTestCase(
+ [
+ SkillTestCasePyFile(
+ input="""
+def function_to_trace():
+ Y()
+ Z()
+
+def Y():
+ A()
+
+def Z():
+ B()
+
+def A():
+ pass
+
+def B():
+ C()
+
+def C():
+ pass
+""",
+ filepath="example.py",
+ )
+ ],
+ graph=True,
+)
+
+
+@skill(eval_skill=False, prompt="Show me a visualization of the call graph from X", uid="81e8fbb7-a00a-4e74-b9c2-24f79d24d389")
+class CallGraphFromNode(Skill, ABC):
+ """This skill creates a directed call graph for a given function. Starting from the specified function, it recursively iterates
+ through its function calls and the functions called by them, building a graph of the call paths to a maximum depth. The root of the directed graph
+ is the starting function, each node represents a function call, and edge from node A to node B indicates that function A calls function B. In its current form,
+ it ignores recursive calls and external modules but can be modified trivially to include them. Furthermore, this skill can easily be adapted to support
+ creating a call graph for a class method. In order to do this one simply needs to replace
+
+ `function_to_trace = codebase.get_function("function_to_trace")`
+
+ with
+
+ `function_to_trace = codebase.get_class("class_of_method_to_trace").get_method("method_to_trace")`
+ """
+
+ @staticmethod
+ @skill_impl(test_cases=[CallGraphFromNodeTest], language=ProgrammingLanguage.PYTHON)
+ @skill_impl(test_cases=[], skip_test=True, language=ProgrammingLanguage.TYPESCRIPT)
+ def skill_func(codebase: CodebaseType):
+ # Create a directed graph
+ G = nx.DiGraph()
+
+ # ===== [ Whether to Graph External Modules] =====
+ GRAPH_EXERNAL_MODULE_CALLS = False
+
+ # ===== [ Maximum Recursive Depth ] =====
+ MAX_DEPTH = 5
+
+ def create_downstream_call_trace(parent: FunctionCall | Function | None = None, depth: int = 0):
+ """Creates call graph for parent
+
+ This function recurses through the call graph of a function and creates a visualization
+
+ Args:
+ parent (FunctionCallDefinition| Function): The function for which a call graph will be created.
+ depth (int): The current depth of the recursive stack.
+
+ """
+ # if the maximum recursive depth has been exceeded return
+ if MAX_DEPTH <= depth:
+ return
+ if isinstance(parent, FunctionCall):
+ src_call, src_func = parent, parent.function_definition
+ else:
+ src_call, src_func = parent, parent
+ # Iterate over all call paths of the symbol
+ for call in src_func.function_calls:
+ # the symbol being called
+ func = call.function_definition
+
+ # ignore direct recursive calls
+ if func.name == src_func.name:
+ continue
+
+ # if the function being called is not from an external module
+ if not isinstance(func, ExternalModule):
+ # add `call` to the graph and an edge from `src_call` to `call`
+ G.add_node(call)
+ G.add_edge(src_call, call)
+
+ # recursive call to function call
+ create_downstream_call_trace(call, depth + 1)
+ elif GRAPH_EXERNAL_MODULE_CALLS:
+ # add `call` to the graph and an edge from `src_call` to `call`
+ G.add_node(call)
+ G.add_edge(src_call, call)
+
+ # ===== [ Function To Be Traced] =====
+ function_to_trace = codebase.get_function("function_to_trace")
+
+ # Set starting node
+ G.add_node(function_to_trace, color="yellow")
+
+ # Add all the children (and sub-children) to the graph
+ create_downstream_call_trace(function_to_trace)
+
+ # Visualize the graph
+ codebase.visualize(G)
+
+
+CallGraphFilterTest = SkillTestCase(
+ [
+ SkillTestCasePyFile(
+ input="""
+class MyClass:
+ def get(self):
+ self.helper_method()
+ return "GET request"
+
+ def post(self):
+ self.helper_method()
+ return "POST request"
+
+ def patch(self):
+ return "PATCH request"
+
+ def delete(self):
+ return "DELETE request"
+
+ def helper_method(self):
+ pass
+
+ def other_method(self):
+ self.helper_method()
+ return "This method should not be included"
+
+def external_function():
+ instance = MyClass()
+ instance.get()
+ instance.post()
+ instance.other_method()
+""",
+ filepath="path/to/file.py",
+ ),
+ SkillTestCasePyFile(
+ input="""
+from path.to.file import MyClass
+
+def function_to_trace():
+ instance = MyClass()
+ assert instance.get() == "GET request"
+ assert instance.post() == "POST request"
+ assert instance.patch() == "PATCH request"
+ assert instance.delete() == "DELETE request"
+""",
+ filepath="path/to/file1.py",
+ ),
+ ],
+ graph=True,
+)
+
+
+@skill(
+ eval_skill=False,
+ prompt="Show me a visualization of the call graph from MyClass and filter out test files and include only the methods that have the name post, get, patch, delete",
+ uid="fc1f3ea0-46e7-460a-88ad-5312d4ca1a12",
+)
+class CallGraphFilter(Skill, ABC):
+ """This skill shows a visualization of the call graph from a given function or symbol.
+ It iterates through the usages of the starting function and its subsequent calls,
+ creating a directed graph of function calls. The skill filters out test files and class declarations
+ and includes only methods with specific names (post, get, patch, delete).
+ The call graph uses red for the starting node, yellow for class methods,
+ and can be customized based on user requests. The graph is limited to a specified depth
+ to manage complexity. In its current form, it ignores recursive calls and external modules
+ but can be modified trivially to include them
+ """
+
+ @staticmethod
+ @skill_impl(test_cases=[CallGraphFilterTest], language=ProgrammingLanguage.PYTHON)
+ @skill_impl(test_cases=[], skip_test=True, language=ProgrammingLanguage.TYPESCRIPT)
+ def skill_func(codebase: CodebaseType):
+ # Create a directed graph
+ G = nx.DiGraph()
+
+ # Get the symbol for my_class
+ func_to_trace = codebase.get_function("function_to_trace")
+
+ # Add the main symbol as a node
+ G.add_node(func_to_trace, color="red")
+
+ # ===== [ Maximum Recursive Depth ] =====
+ MAX_DEPTH = 5
+
+ SKIP_CLASS_DECLARATIONS = True
+
+ cls = codebase.get_class("MyClass")
+
+ # Define a recursive function to traverse function calls
+ def create_filtered_downstream_call_trace(parent: FunctionCall | Function, current_depth, max_depth):
+ if current_depth > max_depth:
+ return
+
+ # if parent is of type Function
+ if isinstance(parent, Function):
+ # set both src_call, src_func to parent
+ src_call, src_func = parent, parent
+ else:
+ # get the first callable of parent
+ src_call, src_func = parent, parent.function_definition
+
+ # Iterate over all call paths of the symbol
+ for call in src_func.function_calls:
+ # the symbol being called
+ func = call.function_definition
+
+ if SKIP_CLASS_DECLARATIONS and isinstance(func, Class):
+ continue
+
+ # if the function being called is not from an external module and is not defined in a test file
+ if not isinstance(func, ExternalModule) and not func.file.filepath.startswith("test"):
+ # add `call` to the graph and an edge from `src_call` to `call`
+ metadata = {}
+ if isinstance(func, Function) and func.is_method and func.name in ["post", "get", "patch", "delete"]:
+ name = f"{func.parent_class.name}.{func.name}"
+ metadata = {"color": "yellow", "name": name}
+ G.add_node(call, **metadata)
+ G.add_edge(src_call, call, symbol=cls) # Add edge from current to successor
+
+ # Recursively add successors of the current symbol
+ create_filtered_downstream_call_trace(call, current_depth + 1, max_depth)
+
+ # Start the recursive traversal
+ create_filtered_downstream_call_trace(func_to_trace, 1, MAX_DEPTH)
+
+ # Visualize the graph
+ codebase.visualize(G)
+
+
+CallPathsBetweenNodesTest = SkillTestCase(
+ [
+ SkillTestCasePyFile(
+ input="""
+def start_func():
+ intermediate_func()
+def intermediate_func():
+ end_func()
+
+def end_func():
+ pass
+""",
+ filepath="example.py",
+ )
+ ],
+ graph=True,
+)
+
+
+@skill(eval_skill=False, prompt="Show me a visualization of the call paths between start_class and end_class", uid="aa3f70c3-ac1c-4737-a8b8-7ba89e3c5671")
+class CallPathsBetweenNodes(Skill, ABC):
+ """This skill generates and visualizes a call graph between two specified functions.
+ It starts from a given function and iteratively traverses through its function calls,
+ building a directed graph of the call paths. The skill then identifies all simple paths between the
+ start and end functions, creating a subgraph that includes only the nodes in these paths.
+
+ By default, the call graph uses blue for the starting node and red for the ending node, but these
+ colors can be customized based on user preferences. The visualization provides a clear representation
+ of how functions are interconnected, helping developers understand the flow of execution and
+ dependencies between different parts of the codebase.
+
+ In its current form, it ignores recursive calls and external modules but can be modified trivially to include them
+ """
+
+ @staticmethod
+ @skill_impl(test_cases=[CallPathsBetweenNodesTest], language=ProgrammingLanguage.PYTHON)
+ @skill_impl(test_cases=[], skip_test=True, language=ProgrammingLanguage.TYPESCRIPT)
+ def skill_func(codebase: CodebaseType):
+ # Create a directed graph
+ G = nx.DiGraph()
+
+ # ===== [ Maximum Recursive Depth ] =====
+ MAX_DEPTH = 5
+
+ # Define a recursive function to traverse usages
+ def create_downstream_call_trace(parent: FunctionCall | Function, end: Callable, current_depth, max_depth):
+ if current_depth > max_depth:
+ return
+
+ # if parent is of type Function
+ if isinstance(parent, Function):
+ # set both src_call, src_func to parent
+ src_call, src_func = parent, parent
+ else:
+ # get the first callable of parent
+ src_call, src_func = parent, parent.function_definition
+
+ # Iterate over all call paths of the symbol
+ for call in src_func.function_calls:
+ # the symbol being called
+ func = call.function_definition
+
+ # ignore direct recursive calls
+ if func.name == src_func.name:
+ continue
+
+ # if the function being called is not from an external module
+ if not isinstance(func, ExternalModule):
+ # add `call` to the graph and an edge from `src_call` to `call`
+ G.add_node(call)
+ G.add_edge(src_call, call)
+
+ if func == end:
+ G.add_edge(call, end)
+ return
+ # recursive call to function call
+ create_downstream_call_trace(call, end, current_depth + 1, max_depth)
+
+ # Get the start and end function
+ start = codebase.get_function("start_func")
+ end = codebase.get_function("end_func")
+
+ # Set starting node as blue
+ G.add_node(start, color="blue")
+ # Set ending node as red
+ G.add_node(end, color="red")
+
+ # Start the recursive traversal
+ create_downstream_call_trace(start, end, 1, MAX_DEPTH)
+
+ # Find all the simple paths between start and end
+ all_paths = nx.all_simple_paths(G, source=start, target=end)
+
+ # Collect all nodes that are part of these paths
+ nodes_in_paths = set()
+ for path in all_paths:
+ nodes_in_paths.update(path)
+
+ # Create a new subgraph with only the nodes in the paths
+ G = G.subgraph(nodes_in_paths)
+
+ # Visualize the graph
+ codebase.visualize(G)
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/method_relationships.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/method_relationships.py
new file mode 100644
index 000000000..b45e1e3fd
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/method_relationships.py
@@ -0,0 +1,107 @@
+import codegen
+import networkx as nx
+from codegen import Codebase
+from codegen.sdk.core.class_definition import Class
+from codegen.sdk.core.detached_symbols.function_call import FunctionCall
+from codegen.sdk.core.external_module import ExternalModule
+from codegen.sdk.core.function import Function
+
+G = nx.DiGraph()
+
+# Configuration Settings
+IGNORE_EXTERNAL_MODULE_CALLS = False
+IGNORE_CLASS_CALLS = True
+MAX_DEPTH = 100
+
+# Track visited nodes to prevent duplicate processing
+visited = set()
+
+COLOR_PALETTE = {
+ "StartMethod": "#9cdcfe", # Light blue for root/entry point methods
+ "PyFunction": "#a277ff", # Purple for regular Python functions
+ "PyClass": "#ffca85", # Warm peach for class definitions
+ "ExternalModule": "#f694ff", # Pink for external module calls
+ "StartClass": "#FFE082", # Yellow for the starting class
+}
+
+
+def graph_class_methods(target_class: Class):
+ """Creates a graph visualization of all methods in a class and their call relationships"""
+ G.add_node(target_class, color=COLOR_PALETTE["StartClass"])
+
+ for method in target_class.methods:
+ method_name = f"{target_class.name}.{method.name}"
+ G.add_node(method, name=method_name, color=COLOR_PALETTE["StartMethod"])
+ visited.add(method)
+ G.add_edge(target_class, method)
+
+ for method in target_class.methods:
+ create_downstream_call_trace(method)
+
+
+def generate_edge_meta(call: FunctionCall) -> dict:
+ """Generate metadata for graph edges representing function calls"""
+ return {"name": call.name, "file_path": call.filepath, "start_point": call.start_point, "end_point": call.end_point, "symbol_name": "FunctionCall"}
+
+
+def create_downstream_call_trace(src_func: Function, depth: int = 0):
+ """Creates call graph for parent function by recursively traversing all function calls"""
+ if MAX_DEPTH <= depth or isinstance(src_func, ExternalModule):
+ return
+
+ for call in src_func.function_calls:
+ if call.name == src_func.name:
+ continue
+
+ func = call.function_definition
+ if not func:
+ continue
+
+ if isinstance(func, ExternalModule) and IGNORE_EXTERNAL_MODULE_CALLS:
+ continue
+ if isinstance(func, Class) and IGNORE_CLASS_CALLS:
+ continue
+
+ if isinstance(func, (Class, ExternalModule)):
+ func_name = func.name
+ elif isinstance(func, Function):
+ func_name = f"{func.parent_class.name}.{func.name}" if func.is_method else func.name
+
+ if func not in visited:
+ G.add_node(func, name=func_name, color=COLOR_PALETTE.get(func.__class__.__name__, None))
+ visited.add(func)
+
+ G.add_edge(src_func, func, **generate_edge_meta(call))
+
+ if isinstance(func, Function):
+ create_downstream_call_trace(func, depth + 1)
+
+
+@codegen.function("visualize-class-method-relationships")
+def run(codebase: Codebase):
+ """Generate a visualization of method call relationships within a class.
+
+ This codemod:
+ 1. Creates a directed graph with the target class as the root node
+ 2. Adds all class methods and their downstream function calls
+ 3. Generates a visual representation of the call hierarchy
+ """
+ global G, visited
+ G = nx.DiGraph()
+ visited = set()
+
+ target_class = codebase.get_class("_Client")
+ graph_class_methods(target_class)
+
+ print(G)
+ print("Use codegen.sh to visualize the graph!")
+
+
+if __name__ == "__main__":
+ print("Initializing codebase...")
+ codebase = Codebase.from_repo("codegen-oss/modal-client", commit="00bf226a1526f9d775d2d70fc7711406aaf42958", language="python")
+ print(f"Codebase with {len(codebase.files)} files and {len(codebase.functions)} functions.")
+ print("Creating graph...")
+
+ run(codebase)
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/viz_cal_graph.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/viz_cal_graph.py
new file mode 100644
index 000000000..095e5f92b
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/call_graph/viz_cal_graph.py
@@ -0,0 +1,121 @@
+import codegen
+import networkx as nx
+from codegen import Codebase
+from codegen.sdk.core.class_definition import Class
+from codegen.sdk.core.detached_symbols.function_call import FunctionCall
+from codegen.sdk.core.external_module import ExternalModule
+from codegen.sdk.core.function import Function
+
+G = nx.DiGraph()
+
+IGNORE_EXTERNAL_MODULE_CALLS = True
+IGNORE_CLASS_CALLS = False
+MAX_DEPTH = 10
+
+# Color scheme for different types of nodes in the visualization
+# Each node type has a distinct color for better visual differentiation
+COLOR_PALETTE = {
+ "StartFunction": "#9cdcfe", # Base purple - draws attention to the root node
+ "PyFunction": "#a277ff", # Mint green - complementary to purple
+ "PyClass": "#ffca85", # Warm peach - provides contrast
+ "ExternalModule": "#f694ff", # Light pink - analogous to base purple
+}
+
+
+def generate_edge_meta(call: FunctionCall) -> dict:
+ """Generate metadata for graph edges representing function calls
+
+ Args:
+ call (FunctionCall): Object containing information about the function call
+
+ Returns:
+ dict: Metadata including name, file path, and location information
+ """
+ return {"name": call.name, "file_path": call.filepath, "start_point": call.start_point, "end_point": call.end_point, "symbol_name": "FunctionCall"}
+
+
+def create_downstream_call_trace(src_func: Function, depth: int = 0):
+ """Creates call graph for parent function by recursively traversing all function calls
+
+ This function builds a directed graph showing all downstream function calls,
+ up to MAX_DEPTH levels deep. Each node represents a function and edges
+ represent calls between functions.
+
+ Args:
+ src_func (Function): The function for which a call graph will be created
+ depth (int): Current depth in the recursive traversal
+ """
+ # Stop recursion if max depth reached
+ if MAX_DEPTH <= depth:
+ return
+ # Stop if the source is an external module
+ if isinstance(src_func, ExternalModule):
+ return
+
+ # Examine each function call made by the source function
+ for call in src_func.function_calls:
+ # Skip recursive calls
+ if call.name == src_func.name:
+ continue
+
+ # Get the function definition being called
+ func = call.function_definition
+
+ # Skip if function definition not found
+ if not func:
+ continue
+ # Apply filtering based on configuration flags
+ if isinstance(func, ExternalModule) and IGNORE_EXTERNAL_MODULE_CALLS:
+ continue
+ if isinstance(func, Class) and IGNORE_CLASS_CALLS:
+ continue
+
+ # Generate the display name for the function
+ # For methods, include the class name
+ if isinstance(func, (Class, ExternalModule)):
+ func_name = func.name
+ elif isinstance(func, Function):
+ func_name = f"{func.parent_class.name}.{func.name}" if func.is_method else func.name
+
+ # Add node and edge to the graph with appropriate metadata
+ G.add_node(func, name=func_name, color=COLOR_PALETTE.get(func.__class__.__name__))
+ G.add_edge(src_func, func, **generate_edge_meta(call))
+
+ # Recursively process called function if it's a regular function
+ if isinstance(func, Function):
+ create_downstream_call_trace(func, depth + 1)
+
+
+@codegen.function("visualize-function-call-relationships")
+def run(codebase: Codebase):
+ """Generate a visualization of function call relationships in a codebase.
+
+ This codemod:
+ 1. Creates a directed graph of function calls starting from a target method
+ 2. Tracks relationships between functions, classes, and external modules
+ 3. Generates a visual representation of the call hierarchy
+ """
+ global G
+ G = nx.DiGraph()
+
+ target_class = codebase.get_class("SharingConfigurationViewSet")
+ target_method = target_class.get_method("patch")
+
+ # Generate the call graph starting from the target method
+ create_downstream_call_trace(target_method)
+
+ # Add the root node (target method) to the graph
+ G.add_node(target_method, name=f"{target_class.name}.{target_method.name}", color=COLOR_PALETTE.get("StartFunction"))
+
+ print(G)
+ print("Use codegen.sh to visualize the graph!")
+
+
+if __name__ == "__main__":
+ print("Initializing codebase...")
+ codebase = Codebase.from_repo("codegen-oss/posthog", commit="b174f2221ea4ae50e715eb6a7e70e9a2b0760800", language="python")
+ print(f"Codebase with {len(codebase.files)} files and {len(codebase.functions)} functions.")
+ print("Creating graph...")
+
+ run(codebase)
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/code_visualizer.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/code_visualizer.py
new file mode 100644
index 000000000..98c462643
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/code_visualizer.py
@@ -0,0 +1,568 @@
+#!/usr/bin/env python3
+"""
+Code Structure Visualizer
+
+This module provides visualization capabilities for code structures such as
+call graphs, dependency graphs, class methods, and blast radius.
+"""
+
+import logging
+
+from .visualizer import BaseVisualizer, OutputFormat, VisualizationType
+
+try:
+ import matplotlib.pyplot as plt
+ import networkx as nx
+except ImportError:
+ logging.warning(
+ "Visualization dependencies not found. Please install them with: pip install networkx matplotlib"
+ )
+
+logger = logging.getLogger(__name__)
+
+
+class CodeVisualizer(BaseVisualizer):
+ """
+ Visualizer for code structures such as call graphs and dependencies.
+
+ This class provides methods to visualize relationships between code entities
+ including functions, classes, and modules.
+ """
+
+ def __init__(self, codebase=None, context=None, **kwargs):
+ """
+ Initialize the CodeVisualizer.
+
+ Args:
+ codebase: Codebase instance to visualize
+ context: Context providing graph representation
+ **kwargs: Additional configuration options
+ """
+ super().__init__(**kwargs)
+ self.codebase = codebase
+ self.context = context
+
+ # Initialize codebase if needed
+ if not self.codebase and not self.context and "analyzer" in kwargs:
+ self.codebase = kwargs["analyzer"].base_codebase
+ self.context = kwargs["analyzer"].base_context
+
+ def visualize_call_graph(self, function_name: str, max_depth: int | None = None):
+ """
+ Generate a call graph visualization for a function.
+
+ Args:
+ function_name: Name of the function to visualize
+ max_depth: Maximum depth of the call graph (overrides config)
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ # Set max depth
+ current_max_depth = (
+ max_depth if max_depth is not None else self.config.max_depth
+ )
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Find the function in the codebase
+ function = None
+ for func in self.codebase.functions:
+ if func.name == function_name:
+ function = func
+ break
+
+ if not function:
+ logger.error(f"Function {function_name} not found in codebase")
+ return None
+
+ # Add root node
+ self._add_node(
+ function,
+ name=function_name,
+ color=self.config.color_palette.get("Root"),
+ is_root=True,
+ )
+
+ # Recursively add call relationships
+ visited = {function}
+
+ def add_calls(func, depth=0):
+ if depth >= current_max_depth:
+ return
+
+ # Skip if no function calls attribute
+ if not hasattr(func, "function_calls"):
+ return
+
+ for call in func.function_calls:
+ # Skip recursive calls
+ if call.name == func.name:
+ continue
+
+ # Get the called function
+ called_func = call.function_definition
+ if not called_func:
+ continue
+
+ # Skip external modules if configured
+ if (
+ self.config.ignore_external
+ and hasattr(called_func, "is_external")
+ and called_func.is_external
+ ):
+ continue
+
+ # Generate name for display
+ if (
+ hasattr(called_func, "is_method")
+ and called_func.is_method
+ and hasattr(called_func, "parent_class")
+ ):
+ called_name = f"{called_func.parent_class.name}.{called_func.name}"
+ else:
+ called_name = called_func.name
+
+ # Add node for called function
+ self._add_node(
+ called_func,
+ name=called_name,
+ color=self.config.color_palette.get("Function"),
+ file_path=called_func.file.path
+ if hasattr(called_func, "file")
+ and hasattr(called_func.file, "path")
+ else None,
+ )
+
+ # Add edge for call relationship
+ self._add_edge(
+ function,
+ called_func,
+ type="call",
+ file_path=call.filepath if hasattr(call, "filepath") else None,
+ line=call.line if hasattr(call, "line") else None,
+ )
+
+ # Recursively process called function
+ if called_func not in visited:
+ visited.add(called_func)
+ add_calls(called_func, depth + 1)
+
+ # Start from the root function
+ add_calls(function)
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.CALL_GRAPH, function_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.CALL_GRAPH, function_name, fig
+ )
+
+ def visualize_dependency_graph(
+ self, symbol_name: str, max_depth: int | None = None
+ ):
+ """
+ Generate a dependency graph visualization for a symbol.
+
+ Args:
+ symbol_name: Name of the symbol to visualize
+ max_depth: Maximum depth of the dependency graph (overrides config)
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ # Set max depth
+ current_max_depth = (
+ max_depth if max_depth is not None else self.config.max_depth
+ )
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Find the symbol in the codebase
+ symbol = None
+ for sym in self.codebase.symbols:
+ if hasattr(sym, "name") and sym.name == symbol_name:
+ symbol = sym
+ break
+
+ if not symbol:
+ logger.error(f"Symbol {symbol_name} not found in codebase")
+ return None
+
+ # Add root node
+ self._add_node(
+ symbol,
+ name=symbol_name,
+ color=self.config.color_palette.get("Root"),
+ is_root=True,
+ )
+
+ # Recursively add dependencies
+ visited = {symbol}
+
+ def add_dependencies(sym, depth=0):
+ if depth >= current_max_depth:
+ return
+
+ # Skip if no dependencies attribute
+ if not hasattr(sym, "dependencies"):
+ return
+
+ for dep in sym.dependencies:
+ dep_symbol = None
+
+ if hasattr(dep, "__class__") and dep.__class__.__name__ == "Symbol":
+ dep_symbol = dep
+ elif hasattr(dep, "resolved_symbol"):
+ dep_symbol = dep.resolved_symbol
+
+ if not dep_symbol:
+ continue
+
+ # Skip external modules if configured
+ if (
+ self.config.ignore_external
+ and hasattr(dep_symbol, "is_external")
+ and dep_symbol.is_external
+ ):
+ continue
+
+ # Add node for dependency
+ self._add_node(
+ dep_symbol,
+ name=dep_symbol.name
+ if hasattr(dep_symbol, "name")
+ else str(dep_symbol),
+ color=self.config.color_palette.get(
+ dep_symbol.__class__.__name__, "#BBBBBB"
+ ),
+ file_path=dep_symbol.file.path
+ if hasattr(dep_symbol, "file") and hasattr(dep_symbol.file, "path")
+ else None,
+ )
+
+ # Add edge for dependency relationship
+ self._add_edge(sym, dep_symbol, type="depends_on")
+
+ # Recursively process dependency
+ if dep_symbol not in visited:
+ visited.add(dep_symbol)
+ add_dependencies(dep_symbol, depth + 1)
+
+ # Start from the root symbol
+ add_dependencies(symbol)
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.DEPENDENCY_GRAPH, symbol_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.DEPENDENCY_GRAPH, symbol_name, fig
+ )
+
+ def visualize_blast_radius(self, symbol_name: str, max_depth: int | None = None):
+ """
+ Generate a blast radius visualization for a symbol.
+
+ Args:
+ symbol_name: Name of the symbol to visualize
+ max_depth: Maximum depth of the blast radius (overrides config)
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ # Set max depth
+ current_max_depth = (
+ max_depth if max_depth is not None else self.config.max_depth
+ )
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Find the symbol in the codebase
+ symbol = None
+ for sym in self.codebase.symbols:
+ if hasattr(sym, "name") and sym.name == symbol_name:
+ symbol = sym
+ break
+
+ if not symbol:
+ logger.error(f"Symbol {symbol_name} not found in codebase")
+ return None
+
+ # Add root node
+ self._add_node(
+ symbol,
+ name=symbol_name,
+ color=self.config.color_palette.get("Root"),
+ is_root=True,
+ )
+
+ # Recursively add usages (reverse dependencies)
+ visited = {symbol}
+
+ def add_usages(sym, depth=0):
+ if depth >= current_max_depth:
+ return
+
+ # Skip if no usages attribute
+ if not hasattr(sym, "usages"):
+ return
+
+ for usage in sym.usages:
+ # Skip if no usage symbol
+ if not hasattr(usage, "usage_symbol"):
+ continue
+
+ usage_symbol = usage.usage_symbol
+
+ # Skip external modules if configured
+ if (
+ self.config.ignore_external
+ and hasattr(usage_symbol, "is_external")
+ and usage_symbol.is_external
+ ):
+ continue
+
+ # Add node for usage
+ self._add_node(
+ usage_symbol,
+ name=usage_symbol.name
+ if hasattr(usage_symbol, "name")
+ else str(usage_symbol),
+ color=self.config.color_palette.get(
+ usage_symbol.__class__.__name__, "#BBBBBB"
+ ),
+ file_path=usage_symbol.file.path
+ if hasattr(usage_symbol, "file")
+ and hasattr(usage_symbol.file, "path")
+ else None,
+ )
+
+ # Add edge for usage relationship
+ self._add_edge(sym, usage_symbol, type="used_by")
+
+ # Recursively process usage
+ if usage_symbol not in visited:
+ visited.add(usage_symbol)
+ add_usages(usage_symbol, depth + 1)
+
+ # Start from the root symbol
+ add_usages(symbol)
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.BLAST_RADIUS, symbol_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.BLAST_RADIUS, symbol_name, fig
+ )
+
+ def visualize_class_methods(self, class_name: str):
+ """
+ Generate a class methods visualization.
+
+ Args:
+ class_name: Name of the class to visualize
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ # Initialize graph
+ self._initialize_graph()
+
+ # Find the class in the codebase
+ class_obj = None
+ for cls in self.codebase.classes:
+ if cls.name == class_name:
+ class_obj = cls
+ break
+
+ if not class_obj:
+ logger.error(f"Class {class_name} not found in codebase")
+ return None
+
+ # Add class node
+ self._add_node(
+ class_obj,
+ name=class_name,
+ color=self.config.color_palette.get("Class"),
+ is_root=True,
+ )
+
+ # Skip if no methods attribute
+ if not hasattr(class_obj, "methods"):
+ logger.error(f"Class {class_name} has no methods attribute")
+ return None
+
+ # Add method nodes and connections
+ method_ids = {}
+ for method in class_obj.methods:
+ method_name = f"{class_name}.{method.name}"
+
+ # Add method node
+ method_id = self._add_node(
+ method,
+ name=method_name,
+ color=self.config.color_palette.get("Function"),
+ file_path=method.file.path
+ if hasattr(method, "file") and hasattr(method.file, "path")
+ else None,
+ )
+
+ method_ids[method.name] = method_id
+
+ # Add edge from class to method
+ self._add_edge(class_obj, method, type="contains")
+
+ # Add call relationships between methods
+ for method in class_obj.methods:
+ # Skip if no function calls attribute
+ if not hasattr(method, "function_calls"):
+ continue
+
+ for call in method.function_calls:
+ # Get the called function
+ called_func = call.function_definition
+ if not called_func:
+ continue
+
+ # Only add edges between methods of this class
+ if (
+ hasattr(called_func, "is_method")
+ and called_func.is_method
+ and hasattr(called_func, "parent_class")
+ and called_func.parent_class == class_obj
+ ):
+ self._add_edge(
+ method,
+ called_func,
+ type="calls",
+ line=call.line if hasattr(call, "line") else None,
+ )
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.CLASS_METHODS, class_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.CLASS_METHODS, class_name, fig
+ )
+
+ def visualize_module_dependencies(self, module_path: str):
+ """
+ Generate a module dependencies visualization.
+
+ Args:
+ module_path: Path to the module to visualize
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ # Initialize graph
+ self._initialize_graph()
+
+ # Get all files in the module
+ module_files = []
+ for file in self.codebase.files:
+ if hasattr(file, "path") and str(file.path).startswith(module_path):
+ module_files.append(file)
+
+ if not module_files:
+ logger.error(f"No files found in module {module_path}")
+ return None
+
+ # Add file nodes
+ module_node_ids = {}
+ for file in module_files:
+ file_name = str(file.path).split("/")[-1]
+ file_module = "/".join(str(file.path).split("/")[:-1])
+
+ # Add file node
+ file_id = self._add_node(
+ file,
+ name=file_name,
+ module=file_module,
+ color=self.config.color_palette.get("File"),
+ file_path=str(file.path),
+ )
+
+ module_node_ids[str(file.path)] = file_id
+
+ # Add import relationships
+ for file in module_files:
+ # Skip if no imports attribute
+ if not hasattr(file, "imports"):
+ continue
+
+ for imp in file.imports:
+ imported_file = None
+
+ # Try to get imported file
+ if hasattr(imp, "resolved_file"):
+ imported_file = imp.resolved_file
+ elif hasattr(imp, "resolved_symbol") and hasattr(
+ imp.resolved_symbol, "file"
+ ):
+ imported_file = imp.resolved_symbol.file
+
+ if not imported_file:
+ continue
+
+ # Skip external modules if configured
+ if (
+ self.config.ignore_external
+ and hasattr(imported_file, "is_external")
+ and imported_file.is_external
+ ):
+ continue
+
+ # Add node for imported file if not already added
+ imported_path = (
+ str(imported_file.path) if hasattr(imported_file, "path") else ""
+ )
+
+ if imported_path not in module_node_ids:
+ imported_name = imported_path.split("/")[-1]
+ imported_module = "/".join(imported_path.split("/")[:-1])
+
+ imported_id = self._add_node(
+ imported_file,
+ name=imported_name,
+ module=imported_module,
+ color=self.config.color_palette.get(
+ "External"
+ if imported_path.startswith(module_path)
+ else "File"
+ ),
+ file_path=imported_path,
+ )
+
+ module_node_ids[imported_path] = imported_id
+
+ # Add edge for import relationship
+ self._add_edge(
+ file,
+ imported_file,
+ type="imports",
+ import_name=imp.name if hasattr(imp, "name") else "",
+ )
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/codebase_visualizer.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/codebase_visualizer.py
new file mode 100644
index 000000000..2cea2331b
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/codebase_visualizer.py
@@ -0,0 +1,1690 @@
+#!/usr/bin/env python3
+"""
+Codebase Visualizer Module
+
+This module provides comprehensive visualization capabilities for codebases and PR analyses.
+It integrates with codebase_analyzer.py and context_codebase.py to provide visual representations
+of code structure, dependencies, and issues. It supports multiple visualization types to help
+developers understand codebase architecture and identify potential problems.
+"""
+
+import json
+import logging
+import os
+import sys
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+try:
+ import matplotlib.pyplot as plt
+ import networkx as nx
+ from matplotlib.colors import LinearSegmentedColormap
+except ImportError:
+ print(
+ "Visualization dependencies not found. Please install them with: pip install networkx matplotlib"
+ )
+ sys.exit(1)
+
+try:
+ from codegen.sdk.core.class_definition import Class
+ from codegen.sdk.core.codebase import Codebase
+ from codegen.sdk.core.detached_symbols.function_call import FunctionCall
+ from codegen.sdk.core.file import SourceFile
+ from codegen.sdk.core.function import Function
+ from codegen.sdk.core.import_resolution import Import
+ from codegen.sdk.core.symbol import Symbol
+ from codegen.sdk.enums import EdgeType, SymbolType
+
+ from codegen_on_oss.codebase_analyzer import (
+ AnalysisType,
+ CodebaseAnalyzer,
+ Issue,
+ IssueSeverity,
+ )
+
+ # Import custom modules
+ from codegen_on_oss.context_codebase import (
+ GLOBAL_FILE_IGNORE_LIST,
+ CodebaseContext,
+ get_node_classes,
+ )
+ from codegen_on_oss.current_code_codebase import get_selected_codebase
+except ImportError:
+ print(
+ "Codegen SDK or custom modules not found. Please ensure all dependencies are installed."
+ )
+ sys.exit(1)
+
+# Configure logging
+logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+ handlers=[logging.StreamHandler()],
+)
+logger = logging.getLogger(__name__)
+
+
+class VisualizationType(str, Enum):
+ """Types of visualizations supported by this module."""
+
+ CALL_GRAPH = "call_graph"
+ DEPENDENCY_GRAPH = "dependency_graph"
+ BLAST_RADIUS = "blast_radius"
+ CLASS_METHODS = "class_methods"
+ MODULE_DEPENDENCIES = "module_dependencies"
+ DEAD_CODE = "dead_code"
+ CYCLOMATIC_COMPLEXITY = "cyclomatic_complexity"
+ ISSUES_HEATMAP = "issues_heatmap"
+ PR_COMPARISON = "pr_comparison"
+
+
+class OutputFormat(str, Enum):
+ """Output formats for visualizations."""
+
+ JSON = "json"
+ PNG = "png"
+ SVG = "svg"
+ HTML = "html"
+ DOT = "dot"
+
+
+@dataclass
+class VisualizationConfig:
+ """Configuration for visualization generation."""
+
+ max_depth: int = 5
+ ignore_external: bool = True
+ ignore_tests: bool = True
+ node_size_base: int = 300
+ edge_width_base: float = 1.0
+ filename_filter: list[str] | None = None
+ symbol_filter: list[str] | None = None
+ output_format: OutputFormat = OutputFormat.JSON
+ output_directory: str | None = None
+ layout_algorithm: str = "spring"
+ highlight_nodes: list[str] = field(default_factory=list)
+ highlight_color: str = "#ff5555"
+ color_palette: dict[str, str] = field(
+ default_factory=lambda: {
+ "Function": "#a277ff", # Purple
+ "Class": "#ffca85", # Orange
+ "File": "#80CBC4", # Teal
+ "Module": "#81D4FA", # Light Blue
+ "Variable": "#B39DDB", # Light Purple
+ "Root": "#ef5350", # Red
+ "Warning": "#FFCA28", # Amber
+ "Error": "#EF5350", # Red
+ "Dead": "#78909C", # Gray
+ "External": "#B0BEC5", # Light Gray
+ }
+ )
+
+
+class CodebaseVisualizer:
+ """
+ Visualizer for codebase structures and analytics.
+
+ This class provides methods to generate various visualizations of a codebase,
+ including call graphs, dependency graphs, complexity heatmaps, and more.
+ It integrates with CodebaseAnalyzer to visualize analysis results.
+ """
+
+ def __init__(
+ self,
+ analyzer: CodebaseAnalyzer | None = None,
+ codebase: Codebase | None = None,
+ context: CodebaseContext | None = None,
+ config: VisualizationConfig | None = None,
+ ):
+ """
+ Initialize the CodebaseVisualizer.
+
+ Args:
+ analyzer: Optional CodebaseAnalyzer instance with analysis results
+ codebase: Optional Codebase instance to visualize
+ context: Optional CodebaseContext providing graph representation
+ config: Visualization configuration options
+ """
+ self.analyzer = analyzer
+ self.codebase = codebase or (analyzer.base_codebase if analyzer else None)
+ self.context = context or (analyzer.base_context if analyzer else None)
+ self.config = config or VisualizationConfig()
+
+ # Create visualization directory if specified
+ if self.config.output_directory:
+ os.makedirs(self.config.output_directory, exist_ok=True)
+
+ # Initialize graph for visualization
+ self.graph = nx.DiGraph()
+
+ # Initialize codebase if needed
+ if not self.codebase and not self.context:
+ logger.info(
+ "No codebase or context provided, initializing from current directory"
+ )
+ self.codebase = get_selected_codebase()
+ self.context = CodebaseContext(
+ codebase=self.codebase, base_path=os.getcwd()
+ )
+ elif self.codebase and not self.context:
+ logger.info("Creating context from provided codebase")
+ self.context = CodebaseContext(
+ codebase=self.codebase,
+ base_path=os.getcwd()
+ if not hasattr(self.codebase, "base_path")
+ else self.codebase.base_path,
+ )
+
+ def _initialize_graph(self):
+ """Initialize a fresh graph for visualization."""
+ self.graph = nx.DiGraph()
+
+ def _add_node(self, node: Any, **attrs):
+ """
+ Add a node to the visualization graph with attributes.
+
+ Args:
+ node: Node object to add
+ **attrs: Node attributes
+ """
+ # Skip if node already exists
+ if self.graph.has_node(node):
+ return
+
+ # Generate node ID (memory address for unique identification)
+ node_id = id(node)
+
+ # Get node name
+ if "name" in attrs:
+ node_name = attrs["name"]
+ elif hasattr(node, "name"):
+ node_name = node.name
+ elif hasattr(node, "path"):
+ node_name = str(node.path).split("/")[-1]
+ else:
+ node_name = str(node)
+
+ # Determine node type and color
+ node_type = node.__class__.__name__
+ color = attrs.get("color", self.config.color_palette.get(node_type, "#BBBBBB"))
+
+ # Add node with attributes
+ self.graph.add_node(
+ node_id,
+ original_node=node,
+ name=node_name,
+ type=node_type,
+ color=color,
+ **attrs,
+ )
+
+ return node_id
+
+ def _add_edge(self, source: Any, target: Any, **attrs):
+ """
+ Add an edge to the visualization graph with attributes.
+
+ Args:
+ source: Source node
+ target: Target node
+ **attrs: Edge attributes
+ """
+ # Get node IDs
+ source_id = id(source)
+ target_id = id(target)
+
+ # Add edge with attributes
+ self.graph.add_edge(source_id, target_id, **attrs)
+
+ def _generate_filename(
+ self, visualization_type: VisualizationType, entity_name: str
+ ):
+ """
+ Generate a filename for the visualization.
+
+ Args:
+ visualization_type: Type of visualization
+ entity_name: Name of the entity being visualized
+
+ Returns:
+ Generated filename
+ """
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+ sanitized_name = (
+ entity_name.replace("/", "_").replace("\\", "_").replace(".", "_")
+ )
+ return f"{visualization_type.value}_{sanitized_name}_{timestamp}.{self.config.output_format.value}"
+
+ def _save_visualization(
+ self, visualization_type: VisualizationType, entity_name: str, data: Any
+ ):
+ """
+ Save a visualization to file or return it.
+
+ Args:
+ visualization_type: Type of visualization
+ entity_name: Name of the entity being visualized
+ data: Visualization data to save
+
+ Returns:
+ Path to saved file or visualization data
+ """
+ filename = self._generate_filename(visualization_type, entity_name)
+
+ if self.config.output_directory:
+ filepath = os.path.join(self.config.output_directory, filename)
+ else:
+ filepath = filename
+
+ if self.config.output_format == OutputFormat.JSON:
+ with open(filepath, "w") as f:
+ json.dump(data, f, indent=2)
+ elif self.config.output_format in [OutputFormat.PNG, OutputFormat.SVG]:
+ # Save matplotlib figure
+ plt.savefig(
+ filepath, format=self.config.output_format.value, bbox_inches="tight"
+ )
+ plt.close()
+ elif self.config.output_format == OutputFormat.DOT:
+ # Save as DOT file for Graphviz
+ try:
+ from networkx.drawing.nx_agraph import write_dot
+
+ write_dot(self.graph, filepath)
+ except ImportError:
+ logger.exception(
+ "networkx.drawing.nx_agraph not available. Install pygraphviz for DOT format."
+ )
+ return None
+
+ logger.info(f"Visualization saved to {filepath}")
+ return filepath
+
+ def _convert_graph_to_json(self):
+ """
+ Convert the networkx graph to a JSON-serializable dictionary.
+
+ Returns:
+ Dictionary representation of the graph
+ """
+ nodes = []
+ for node, attrs in self.graph.nodes(data=True):
+ # Create a serializable node
+ node_data = {
+ "id": node,
+ "name": attrs.get("name", ""),
+ "type": attrs.get("type", ""),
+ "color": attrs.get("color", "#BBBBBB"),
+ }
+
+ # Add file path if available
+ if "file_path" in attrs:
+ node_data["file_path"] = attrs["file_path"]
+
+ # Add other attributes
+ for key, value in attrs.items():
+ if key not in ["name", "type", "color", "file_path", "original_node"]:
+ if (
+ isinstance(value, str | int | float | bool | list | dict)
+ or value is None
+ ):
+ node_data[key] = value
+
+ nodes.append(node_data)
+
+ edges = []
+ for source, target, attrs in self.graph.edges(data=True):
+ # Create a serializable edge
+ edge_data = {
+ "source": source,
+ "target": target,
+ }
+
+ # Add other attributes
+ for key, value in attrs.items():
+ if (
+ isinstance(value, str | int | float | bool | list | dict)
+ or value is None
+ ):
+ edge_data[key] = value
+
+ edges.append(edge_data)
+
+ return {
+ "nodes": nodes,
+ "edges": edges,
+ "metadata": {
+ "visualization_type": self.current_visualization_type,
+ "entity_name": self.current_entity_name,
+ "timestamp": datetime.now().isoformat(),
+ "node_count": len(nodes),
+ "edge_count": len(edges),
+ },
+ }
+
+ def _plot_graph(self):
+ """
+ Plot the graph using matplotlib.
+
+ Returns:
+ Matplotlib figure
+ """
+ plt.figure(figsize=(12, 10))
+
+ # Extract node positions using specified layout algorithm
+ if self.config.layout_algorithm == "spring":
+ pos = nx.spring_layout(self.graph, seed=42)
+ elif self.config.layout_algorithm == "kamada_kawai":
+ pos = nx.kamada_kawai_layout(self.graph)
+ elif self.config.layout_algorithm == "spectral":
+ pos = nx.spectral_layout(self.graph)
+ else:
+ # Default to spring layout
+ pos = nx.spring_layout(self.graph, seed=42)
+
+ # Extract node colors
+ node_colors = [
+ attrs.get("color", "#BBBBBB") for _, attrs in self.graph.nodes(data=True)
+ ]
+
+ # Extract node sizes (can be based on some metric)
+ node_sizes = [self.config.node_size_base for _ in self.graph.nodes()]
+
+ # Draw nodes
+ nx.draw_networkx_nodes(
+ self.graph, pos, node_color=node_colors, node_size=node_sizes, alpha=0.8
+ )
+
+ # Draw edges
+ nx.draw_networkx_edges(
+ self.graph,
+ pos,
+ width=self.config.edge_width_base,
+ alpha=0.6,
+ arrows=True,
+ arrowsize=10,
+ )
+
+ # Draw labels
+ nx.draw_networkx_labels(
+ self.graph,
+ pos,
+ labels={
+ node: attrs.get("name", "")
+ for node, attrs in self.graph.nodes(data=True)
+ },
+ font_size=8,
+ font_weight="bold",
+ )
+
+ plt.title(f"{self.current_visualization_type} - {self.current_entity_name}")
+ plt.axis("off")
+
+ return plt.gcf()
+
+ def visualize_call_graph(self, function_name: str, max_depth: int | None = None):
+ """
+ Generate a call graph visualization for a function.
+
+ Args:
+ function_name: Name of the function to visualize
+ max_depth: Maximum depth of the call graph (overrides config)
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ self.current_visualization_type = VisualizationType.CALL_GRAPH
+ self.current_entity_name = function_name
+
+ # Set max depth
+ current_max_depth = (
+ max_depth if max_depth is not None else self.config.max_depth
+ )
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Find the function in the codebase
+ function = None
+ for func in self.codebase.functions:
+ if func.name == function_name:
+ function = func
+ break
+
+ if not function:
+ logger.error(f"Function {function_name} not found in codebase")
+ return None
+
+ # Add root node
+ self._add_node(
+ function,
+ name=function_name,
+ color=self.config.color_palette.get("Root"),
+ is_root=True,
+ )
+
+ # Recursively add call relationships
+ visited = {function}
+
+ def add_calls(func, depth=0):
+ if depth >= current_max_depth:
+ return
+
+ # Skip if no function calls attribute
+ if not hasattr(func, "function_calls"):
+ return
+
+ for call in func.function_calls:
+ # Skip recursive calls
+ if call.name == func.name:
+ continue
+
+ # Get the called function
+ called_func = call.function_definition
+ if not called_func:
+ continue
+
+ # Skip external modules if configured
+ if (
+ self.config.ignore_external
+ and hasattr(called_func, "is_external")
+ and called_func.is_external
+ ):
+ continue
+
+ # Generate name for display
+ if (
+ hasattr(called_func, "is_method")
+ and called_func.is_method
+ and hasattr(called_func, "parent_class")
+ ):
+ called_name = f"{called_func.parent_class.name}.{called_func.name}"
+ else:
+ called_name = called_func.name
+
+ # Add node for called function
+ self._add_node(
+ called_func,
+ name=called_name,
+ color=self.config.color_palette.get("Function"),
+ file_path=called_func.file.path
+ if hasattr(called_func, "file")
+ and hasattr(called_func.file, "path")
+ else None,
+ )
+
+ # Add edge for call relationship
+ self._add_edge(
+ function,
+ called_func,
+ type="call",
+ file_path=call.filepath if hasattr(call, "filepath") else None,
+ line=call.line if hasattr(call, "line") else None,
+ )
+
+ # Recursively process called function
+ if isinstance(called_func, Function) and called_func not in visited:
+ visited.add(called_func)
+ add_calls(called_func, depth + 1)
+
+ # Start from the root function
+ add_calls(function)
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.CALL_GRAPH, function_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.CALL_GRAPH, function_name, fig
+ )
+
+ def visualize_dependency_graph(
+ self, symbol_name: str, max_depth: int | None = None
+ ):
+ """
+ Generate a dependency graph visualization for a symbol.
+
+ Args:
+ symbol_name: Name of the symbol to visualize
+ max_depth: Maximum depth of the dependency graph (overrides config)
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ self.current_visualization_type = VisualizationType.DEPENDENCY_GRAPH
+ self.current_entity_name = symbol_name
+
+ # Set max depth
+ current_max_depth = (
+ max_depth if max_depth is not None else self.config.max_depth
+ )
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Find the symbol in the codebase
+ symbol = None
+ for sym in self.codebase.symbols:
+ if hasattr(sym, "name") and sym.name == symbol_name:
+ symbol = sym
+ break
+
+ if not symbol:
+ logger.error(f"Symbol {symbol_name} not found in codebase")
+ return None
+
+ # Add root node
+ self._add_node(
+ symbol,
+ name=symbol_name,
+ color=self.config.color_palette.get("Root"),
+ is_root=True,
+ )
+
+ # Recursively add dependencies
+ visited = {symbol}
+
+ def add_dependencies(sym, depth=0):
+ if depth >= current_max_depth:
+ return
+
+ # Skip if no dependencies attribute
+ if not hasattr(sym, "dependencies"):
+ return
+
+ for dep in sym.dependencies:
+ dep_symbol = None
+
+ if isinstance(dep, Symbol):
+ dep_symbol = dep
+ elif isinstance(dep, Import) and hasattr(dep, "resolved_symbol"):
+ dep_symbol = dep.resolved_symbol
+
+ if not dep_symbol:
+ continue
+
+ # Skip external modules if configured
+ if (
+ self.config.ignore_external
+ and hasattr(dep_symbol, "is_external")
+ and dep_symbol.is_external
+ ):
+ continue
+
+ # Add node for dependency
+ self._add_node(
+ dep_symbol,
+ name=dep_symbol.name
+ if hasattr(dep_symbol, "name")
+ else str(dep_symbol),
+ color=self.config.color_palette.get(
+ dep_symbol.__class__.__name__, "#BBBBBB"
+ ),
+ file_path=dep_symbol.file.path
+ if hasattr(dep_symbol, "file") and hasattr(dep_symbol.file, "path")
+ else None,
+ )
+
+ # Add edge for dependency relationship
+ self._add_edge(sym, dep_symbol, type="depends_on")
+
+ # Recursively process dependency
+ if dep_symbol not in visited:
+ visited.add(dep_symbol)
+ add_dependencies(dep_symbol, depth + 1)
+
+ # Start from the root symbol
+ add_dependencies(symbol)
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.DEPENDENCY_GRAPH, symbol_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.DEPENDENCY_GRAPH, symbol_name, fig
+ )
+
+ def visualize_blast_radius(self, symbol_name: str, max_depth: int | None = None):
+ """
+ Generate a blast radius visualization for a symbol.
+
+ Args:
+ symbol_name: Name of the symbol to visualize
+ max_depth: Maximum depth of the blast radius (overrides config)
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ self.current_visualization_type = VisualizationType.BLAST_RADIUS
+ self.current_entity_name = symbol_name
+
+ # Set max depth
+ current_max_depth = (
+ max_depth if max_depth is not None else self.config.max_depth
+ )
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Find the symbol in the codebase
+ symbol = None
+ for sym in self.codebase.symbols:
+ if hasattr(sym, "name") and sym.name == symbol_name:
+ symbol = sym
+ break
+
+ if not symbol:
+ logger.error(f"Symbol {symbol_name} not found in codebase")
+ return None
+
+ # Add root node
+ self._add_node(
+ symbol,
+ name=symbol_name,
+ color=self.config.color_palette.get("Root"),
+ is_root=True,
+ )
+
+ # Recursively add usages (reverse dependencies)
+ visited = {symbol}
+
+ def add_usages(sym, depth=0):
+ if depth >= current_max_depth:
+ return
+
+ # Skip if no usages attribute
+ if not hasattr(sym, "usages"):
+ return
+
+ for usage in sym.usages:
+ # Skip if no usage symbol
+ if not hasattr(usage, "usage_symbol"):
+ continue
+
+ usage_symbol = usage.usage_symbol
+
+ # Skip external modules if configured
+ if (
+ self.config.ignore_external
+ and hasattr(usage_symbol, "is_external")
+ and usage_symbol.is_external
+ ):
+ continue
+
+ # Add node for usage
+ self._add_node(
+ usage_symbol,
+ name=usage_symbol.name
+ if hasattr(usage_symbol, "name")
+ else str(usage_symbol),
+ color=self.config.color_palette.get(
+ usage_symbol.__class__.__name__, "#BBBBBB"
+ ),
+ file_path=usage_symbol.file.path
+ if hasattr(usage_symbol, "file")
+ and hasattr(usage_symbol.file, "path")
+ else None,
+ )
+
+ # Add edge for usage relationship
+ self._add_edge(sym, usage_symbol, type="used_by")
+
+ # Recursively process usage
+ if usage_symbol not in visited:
+ visited.add(usage_symbol)
+ add_usages(usage_symbol, depth + 1)
+
+ # Start from the root symbol
+ add_usages(symbol)
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.BLAST_RADIUS, symbol_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.BLAST_RADIUS, symbol_name, fig
+ )
+
+ def visualize_class_methods(self, class_name: str):
+ """
+ Generate a class methods visualization.
+
+ Args:
+ class_name: Name of the class to visualize
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ self.current_visualization_type = VisualizationType.CLASS_METHODS
+ self.current_entity_name = class_name
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Find the class in the codebase
+ class_obj = None
+ for cls in self.codebase.classes:
+ if cls.name == class_name:
+ class_obj = cls
+ break
+
+ if not class_obj:
+ logger.error(f"Class {class_name} not found in codebase")
+ return None
+
+ # Add class node
+ self._add_node(
+ class_obj,
+ name=class_name,
+ color=self.config.color_palette.get("Class"),
+ is_root=True,
+ )
+
+ # Skip if no methods attribute
+ if not hasattr(class_obj, "methods"):
+ logger.error(f"Class {class_name} has no methods attribute")
+ return None
+
+ # Add method nodes and connections
+ method_ids = {}
+ for method in class_obj.methods:
+ method_name = f"{class_name}.{method.name}"
+
+ # Add method node
+ method_id = self._add_node(
+ method,
+ name=method_name,
+ color=self.config.color_palette.get("Function"),
+ file_path=method.file.path
+ if hasattr(method, "file") and hasattr(method.file, "path")
+ else None,
+ )
+
+ method_ids[method.name] = method_id
+
+ # Add edge from class to method
+ self._add_edge(class_obj, method, type="contains")
+
+ # Add call relationships between methods
+ for method in class_obj.methods:
+ # Skip if no function calls attribute
+ if not hasattr(method, "function_calls"):
+ continue
+
+ for call in method.function_calls:
+ # Get the called function
+ called_func = call.function_definition
+ if not called_func:
+ continue
+
+ # Only add edges between methods of this class
+ if (
+ hasattr(called_func, "is_method")
+ and called_func.is_method
+ and hasattr(called_func, "parent_class")
+ and called_func.parent_class == class_obj
+ ):
+ self._add_edge(
+ method,
+ called_func,
+ type="calls",
+ line=call.line if hasattr(call, "line") else None,
+ )
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.CLASS_METHODS, class_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.CLASS_METHODS, class_name, fig
+ )
+
+ def visualize_module_dependencies(self, module_path: str):
+ """
+ Generate a module dependencies visualization.
+
+ Args:
+ module_path: Path to the module to visualize
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ self.current_visualization_type = VisualizationType.MODULE_DEPENDENCIES
+ self.current_entity_name = module_path
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Get all files in the module
+ module_files = []
+ for file in self.codebase.files:
+ if hasattr(file, "path") and str(file.path).startswith(module_path):
+ module_files.append(file)
+
+ if not module_files:
+ logger.error(f"No files found in module {module_path}")
+ return None
+
+ # Add file nodes
+ module_node_ids = {}
+ for file in module_files:
+ file_name = str(file.path).split("/")[-1]
+ file_module = "/".join(str(file.path).split("/")[:-1])
+
+ # Add file node
+ file_id = self._add_node(
+ file,
+ name=file_name,
+ module=file_module,
+ color=self.config.color_palette.get("File"),
+ file_path=str(file.path),
+ )
+
+ module_node_ids[str(file.path)] = file_id
+
+ # Add import relationships
+ for file in module_files:
+ # Skip if no imports attribute
+ if not hasattr(file, "imports"):
+ continue
+
+ for imp in file.imports:
+ imported_file = None
+
+ # Try to get imported file
+ if hasattr(imp, "resolved_file"):
+ imported_file = imp.resolved_file
+ elif hasattr(imp, "resolved_symbol") and hasattr(
+ imp.resolved_symbol, "file"
+ ):
+ imported_file = imp.resolved_symbol.file
+
+ if not imported_file:
+ continue
+
+ # Skip external modules if configured
+ if (
+ self.config.ignore_external
+ and hasattr(imported_file, "is_external")
+ and imported_file.is_external
+ ):
+ continue
+
+ # Add node for imported file if not already added
+ imported_path = (
+ str(imported_file.path) if hasattr(imported_file, "path") else ""
+ )
+
+ if imported_path not in module_node_ids:
+ imported_name = imported_path.split("/")[-1]
+ imported_module = "/".join(imported_path.split("/")[:-1])
+
+ imported_id = self._add_node(
+ imported_file,
+ name=imported_name,
+ module=imported_module,
+ color=self.config.color_palette.get(
+ "External"
+ if imported_path.startswith(module_path)
+ else "File"
+ ),
+ file_path=imported_path,
+ )
+
+ module_node_ids[imported_path] = imported_id
+
+ # Add edge for import relationship
+ self._add_edge(
+ file,
+ imported_file,
+ type="imports",
+ import_name=imp.name if hasattr(imp, "name") else "",
+ )
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.MODULE_DEPENDENCIES, module_path, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.MODULE_DEPENDENCIES, module_path, fig
+ )
+
+ def visualize_dead_code(self, path_filter: str | None = None):
+ """
+ Generate a visualization of dead (unused) code in the codebase.
+
+ Args:
+ path_filter: Optional path to filter files
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ self.current_visualization_type = VisualizationType.DEAD_CODE
+ self.current_entity_name = path_filter or "codebase"
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Initialize analyzer if needed
+ if not self.analyzer:
+ logger.info("Initializing analyzer for dead code detection")
+ self.analyzer = CodebaseAnalyzer(
+ codebase=self.codebase,
+ repo_path=self.context.base_path
+ if hasattr(self.context, "base_path")
+ else None,
+ )
+
+ # Perform analysis if not already done
+ if not hasattr(self.analyzer, "results") or not self.analyzer.results:
+ logger.info("Running code analysis")
+ self.analyzer.analyze(AnalysisType.CODEBASE)
+
+ # Extract dead code information from analysis results
+ if not hasattr(self.analyzer, "results"):
+ logger.error("Analysis results not available")
+ return None
+
+ dead_code = {}
+ if (
+ "static_analysis" in self.analyzer.results
+ and "dead_code" in self.analyzer.results["static_analysis"]
+ ):
+ dead_code = self.analyzer.results["static_analysis"]["dead_code"]
+
+ if not dead_code:
+ logger.warning("No dead code detected in analysis results")
+ return None
+
+ # Create file nodes for containing dead code
+ file_nodes = {}
+
+ # Process unused functions
+ if "unused_functions" in dead_code:
+ for unused_func in dead_code["unused_functions"]:
+ file_path = unused_func.get("file", "")
+
+ # Skip if path filter is specified and doesn't match
+ if path_filter and not file_path.startswith(path_filter):
+ continue
+
+ # Add file node if not already added
+ if file_path not in file_nodes:
+ # Find file in codebase
+ file_obj = None
+ for file in self.codebase.files:
+ if hasattr(file, "path") and str(file.path) == file_path:
+ file_obj = file
+ break
+
+ if file_obj:
+ file_name = file_path.split("/")[-1]
+ self._add_node(
+ file_obj,
+ name=file_name,
+ color=self.config.color_palette.get("File"),
+ file_path=file_path,
+ )
+
+ file_nodes[file_path] = file_obj
+
+ # Add unused function node
+ func_name = unused_func.get("name", "")
+ func_line = unused_func.get("line", None)
+
+ # Create a placeholder for the function (we don't have the actual object)
+ func_obj = {
+ "name": func_name,
+ "file_path": file_path,
+ "line": func_line,
+ "type": "Function",
+ }
+
+ self._add_node(
+ func_obj,
+ name=func_name,
+ color=self.config.color_palette.get("Dead"),
+ file_path=file_path,
+ line=func_line,
+ is_dead=True,
+ )
+
+ # Add edge from file to function
+ if file_path in file_nodes:
+ self._add_edge(
+ file_nodes[file_path], func_obj, type="contains_dead"
+ )
+
+ # Process unused variables
+ if "unused_variables" in dead_code:
+ for unused_var in dead_code["unused_variables"]:
+ file_path = unused_var.get("file", "")
+
+ # Skip if path filter is specified and doesn't match
+ if path_filter and not file_path.startswith(path_filter):
+ continue
+
+ # Add file node if not already added
+ if file_path not in file_nodes:
+ # Find file in codebase
+ file_obj = None
+ for file in self.codebase.files:
+ if hasattr(file, "path") and str(file.path) == file_path:
+ file_obj = file
+ break
+
+ if file_obj:
+ file_name = file_path.split("/")[-1]
+ self._add_node(
+ file_obj,
+ name=file_name,
+ color=self.config.color_palette.get("File"),
+ file_path=file_path,
+ )
+
+ file_nodes[file_path] = file_obj
+
+ # Add unused variable node
+ var_name = unused_var.get("name", "")
+ var_line = unused_var.get("line", None)
+
+ # Create a placeholder for the variable
+ var_obj = {
+ "name": var_name,
+ "file_path": file_path,
+ "line": var_line,
+ "type": "Variable",
+ }
+
+ self._add_node(
+ var_obj,
+ name=var_name,
+ color=self.config.color_palette.get("Dead"),
+ file_path=file_path,
+ line=var_line,
+ is_dead=True,
+ )
+
+ # Add edge from file to variable
+ if file_path in file_nodes:
+ self._add_edge(file_nodes[file_path], var_obj, type="contains_dead")
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.DEAD_CODE, self.current_entity_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.DEAD_CODE, self.current_entity_name, fig
+ )
+
+ def visualize_cyclomatic_complexity(self, path_filter: str | None = None):
+ """
+ Generate a heatmap visualization of cyclomatic complexity.
+
+ Args:
+ path_filter: Optional path to filter files
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ self.current_visualization_type = VisualizationType.CYCLOMATIC_COMPLEXITY
+ self.current_entity_name = path_filter or "codebase"
+
+ # Initialize analyzer if needed
+ if not self.analyzer:
+ logger.info("Initializing analyzer for complexity analysis")
+ self.analyzer = CodebaseAnalyzer(
+ codebase=self.codebase,
+ repo_path=self.context.base_path
+ if hasattr(self.context, "base_path")
+ else None,
+ )
+
+ # Perform analysis if not already done
+ if not hasattr(self.analyzer, "results") or not self.analyzer.results:
+ logger.info("Running code analysis")
+ self.analyzer.analyze(AnalysisType.CODEBASE)
+
+ # Extract complexity information from analysis results
+ if not hasattr(self.analyzer, "results"):
+ logger.error("Analysis results not available")
+ return None
+
+ complexity_data = {}
+ if (
+ "static_analysis" in self.analyzer.results
+ and "code_complexity" in self.analyzer.results["static_analysis"]
+ ):
+ complexity_data = self.analyzer.results["static_analysis"][
+ "code_complexity"
+ ]
+
+ if not complexity_data:
+ logger.warning("No complexity data found in analysis results")
+ return None
+
+ # Extract function complexities
+ functions = []
+ if "function_complexity" in complexity_data:
+ for func_data in complexity_data["function_complexity"]:
+ # Skip if path filter is specified and doesn't match
+ if path_filter and not func_data.get("file", "").startswith(
+ path_filter
+ ):
+ continue
+
+ functions.append({
+ "name": func_data.get("name", ""),
+ "file": func_data.get("file", ""),
+ "complexity": func_data.get("complexity", 1),
+ "line": func_data.get("line", None),
+ })
+
+ # Sort functions by complexity (descending)
+ functions.sort(key=lambda x: x.get("complexity", 0), reverse=True)
+
+ # Generate heatmap visualization
+ plt.figure(figsize=(12, 10))
+
+ # Extract data for heatmap
+ func_names = [
+ f"{func['name']} ({func['file'].split('/')[-1]})" for func in functions[:30]
+ ]
+ complexities = [func.get("complexity", 0) for func in functions[:30]]
+
+ # Create horizontal bar chart
+ bars = plt.barh(func_names, complexities)
+
+ # Color bars by complexity
+ norm = plt.Normalize(1, max(10, max(complexities)))
+ cmap = plt.cm.get_cmap("YlOrRd")
+
+ for i, bar in enumerate(bars):
+ complexity = complexities[i]
+ bar.set_color(cmap(norm(complexity)))
+
+ # Add labels and title
+ plt.xlabel("Cyclomatic Complexity")
+ plt.title("Top Functions by Cyclomatic Complexity")
+ plt.grid(axis="x", linestyle="--", alpha=0.6)
+
+ # Add colorbar
+ plt.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=cmap), label="Complexity")
+
+ # Save and return visualization
+ return self._save_visualization(
+ VisualizationType.CYCLOMATIC_COMPLEXITY, self.current_entity_name, plt.gcf()
+ )
+
+ def visualize_issues_heatmap(
+ self,
+ severity: IssueSeverity | None = None,
+ path_filter: str | None = None,
+ ):
+ """
+ Generate a heatmap visualization of issues in the codebase.
+
+ Args:
+ severity: Optional severity level to filter issues
+ path_filter: Optional path to filter files
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ self.current_visualization_type = VisualizationType.ISSUES_HEATMAP
+ self.current_entity_name = f"{severity.value if severity else 'all'}_issues"
+
+ # Initialize analyzer if needed
+ if not self.analyzer:
+ logger.info("Initializing analyzer for issues analysis")
+ self.analyzer = CodebaseAnalyzer(
+ codebase=self.codebase,
+ repo_path=self.context.base_path
+ if hasattr(self.context, "base_path")
+ else None,
+ )
+
+ # Perform analysis if not already done
+ if not hasattr(self.analyzer, "results") or not self.analyzer.results:
+ logger.info("Running code analysis")
+ self.analyzer.analyze(AnalysisType.CODEBASE)
+
+ # Extract issues from analysis results
+ if (
+ not hasattr(self.analyzer, "results")
+ or "issues" not in self.analyzer.results
+ ):
+ logger.error("Issues not available in analysis results")
+ return None
+
+ issues = self.analyzer.results["issues"]
+
+ # Filter issues by severity if specified
+ if severity:
+ issues = [issue for issue in issues if issue.get("severity") == severity]
+
+ # Filter issues by path if specified
+ if path_filter:
+ issues = [
+ issue
+ for issue in issues
+ if issue.get("file", "").startswith(path_filter)
+ ]
+
+ if not issues:
+ logger.warning("No issues found matching the criteria")
+ return None
+
+ # Group issues by file
+ file_issues = {}
+ for issue in issues:
+ file_path = issue.get("file", "")
+ if file_path not in file_issues:
+ file_issues[file_path] = []
+
+ file_issues[file_path].append(issue)
+
+ # Generate heatmap visualization
+ plt.figure(figsize=(12, 10))
+
+ # Extract data for heatmap
+ files = list(file_issues.keys())
+ file_names = [file_path.split("/")[-1] for file_path in files]
+ issue_counts = [len(file_issues[file_path]) for file_path in files]
+
+ # Sort by issue count
+ sorted_data = sorted(
+ zip(file_names, issue_counts, files, strict=False),
+ key=lambda x: x[1],
+ reverse=True,
+ )
+ file_names, issue_counts, files = zip(*sorted_data, strict=False)
+
+ # Create horizontal bar chart
+ bars = plt.barh(file_names[:20], issue_counts[:20])
+
+ # Color bars by issue count
+ norm = plt.Normalize(1, max(5, max(issue_counts[:20])))
+ cmap = plt.cm.get_cmap("OrRd")
+
+ for i, bar in enumerate(bars):
+ count = issue_counts[i]
+ bar.set_color(cmap(norm(count)))
+
+ # Add labels and title
+ plt.xlabel("Number of Issues")
+ severity_text = f" ({severity.value})" if severity else ""
+ plt.title(f"Files with the Most Issues{severity_text}")
+ plt.grid(axis="x", linestyle="--", alpha=0.6)
+
+ # Add colorbar
+ plt.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=cmap), label="Issue Count")
+
+ # Save and return visualization
+ return self._save_visualization(
+ VisualizationType.ISSUES_HEATMAP, self.current_entity_name, plt.gcf()
+ )
+
+ def visualize_pr_comparison(self):
+ """
+ Generate a visualization comparing base branch with PR.
+
+ Returns:
+ Visualization data or path to saved file
+ """
+ self.current_visualization_type = VisualizationType.PR_COMPARISON
+
+ # Check if analyzer has PR data
+ if (
+ not self.analyzer
+ or not self.analyzer.pr_codebase
+ or not self.analyzer.base_codebase
+ ):
+ logger.error("PR comparison requires analyzer with PR data")
+ return None
+
+ self.current_entity_name = (
+ f"pr_{self.analyzer.pr_number}"
+ if self.analyzer.pr_number
+ else "pr_comparison"
+ )
+
+ # Perform comparison analysis if not already done
+ if not hasattr(self.analyzer, "results") or not self.analyzer.results:
+ logger.info("Running PR comparison analysis")
+ self.analyzer.analyze(AnalysisType.COMPARISON)
+
+ # Extract comparison data from analysis results
+ if (
+ not hasattr(self.analyzer, "results")
+ or "comparison" not in self.analyzer.results
+ ):
+ logger.error("Comparison data not available in analysis results")
+ return None
+
+ comparison = self.analyzer.results["comparison"]
+
+ # Initialize graph
+ self._initialize_graph()
+
+ # Process symbol comparison data
+ if "symbol_comparison" in comparison:
+ for symbol_data in comparison["symbol_comparison"]:
+ symbol_name = symbol_data.get("name", "")
+ in_base = symbol_data.get("in_base", False)
+ in_pr = symbol_data.get("in_pr", False)
+
+ # Create a placeholder for the symbol
+ symbol_obj = {
+ "name": symbol_name,
+ "in_base": in_base,
+ "in_pr": in_pr,
+ "type": "Symbol",
+ }
+
+ # Determine node color based on presence in base and PR
+ if in_base and in_pr:
+ color = "#A5D6A7" # Light green (modified)
+ elif in_base:
+ color = "#EF9A9A" # Light red (removed)
+ else:
+ color = "#90CAF9" # Light blue (added)
+
+ # Add node for symbol
+ self._add_node(
+ symbol_obj,
+ name=symbol_name,
+ color=color,
+ in_base=in_base,
+ in_pr=in_pr,
+ )
+
+ # Process parameter changes if available
+ if "parameter_changes" in symbol_data:
+ param_changes = symbol_data["parameter_changes"]
+
+ # Process removed parameters
+ for param in param_changes.get("removed", []):
+ param_obj = {
+ "name": param,
+ "change_type": "removed",
+ "type": "Parameter",
+ }
+
+ self._add_node(
+ param_obj,
+ name=param,
+ color="#EF9A9A", # Light red (removed)
+ change_type="removed",
+ )
+
+ self._add_edge(symbol_obj, param_obj, type="removed_parameter")
+
+ # Process added parameters
+ for param in param_changes.get("added", []):
+ param_obj = {
+ "name": param,
+ "change_type": "added",
+ "type": "Parameter",
+ }
+
+ self._add_node(
+ param_obj,
+ name=param,
+ color="#90CAF9", # Light blue (added)
+ change_type="added",
+ )
+
+ self._add_edge(symbol_obj, param_obj, type="added_parameter")
+
+ # Process return type changes if available
+ if "return_type_change" in symbol_data:
+ return_type_change = symbol_data["return_type_change"]
+ old_type = return_type_change.get("old", "None")
+ new_type = return_type_change.get("new", "None")
+
+ return_obj = {
+ "name": f"{old_type} -> {new_type}",
+ "old_type": old_type,
+ "new_type": new_type,
+ "type": "ReturnType",
+ }
+
+ self._add_node(
+ return_obj,
+ name=f"{old_type} -> {new_type}",
+ color="#FFD54F", # Amber (changed)
+ old_type=old_type,
+ new_type=new_type,
+ )
+
+ self._add_edge(symbol_obj, return_obj, type="return_type_change")
+
+ # Process call site issues if available
+ if "call_site_issues" in symbol_data:
+ for issue in symbol_data["call_site_issues"]:
+ issue_file = issue.get("file", "")
+ issue_line = issue.get("line", None)
+ issue_text = issue.get("issue", "")
+
+ # Create a placeholder for the issue
+ issue_obj = {
+ "name": issue_text,
+ "file": issue_file,
+ "line": issue_line,
+ "type": "Issue",
+ }
+
+ self._add_node(
+ issue_obj,
+ name=f"{issue_file.split('/')[-1]}:{issue_line}",
+ color="#EF5350", # Red (error)
+ file_path=issue_file,
+ line=issue_line,
+ issue_text=issue_text,
+ )
+
+ self._add_edge(symbol_obj, issue_obj, type="call_site_issue")
+
+ # Generate visualization data
+ if self.config.output_format == OutputFormat.JSON:
+ data = self._convert_graph_to_json()
+ return self._save_visualization(
+ VisualizationType.PR_COMPARISON, self.current_entity_name, data
+ )
+ else:
+ fig = self._plot_graph()
+ return self._save_visualization(
+ VisualizationType.PR_COMPARISON, self.current_entity_name, fig
+ )
+
+
+# Command-line interface
+def main():
+ """
+ Command-line interface for the codebase visualizer.
+
+ This function parses command-line arguments and generates visualizations
+ based on the specified parameters.
+ """
+ parser = argparse.ArgumentParser(
+ description="Generate visualizations of codebase structure and analysis."
+ )
+
+ # Repository options
+ repo_group = parser.add_argument_group("Repository Options")
+ repo_group.add_argument("--repo-url", help="URL of the repository to analyze")
+ repo_group.add_argument(
+ "--repo-path", help="Local path to the repository to analyze"
+ )
+ repo_group.add_argument("--language", help="Programming language of the codebase")
+
+ # Visualization options
+ viz_group = parser.add_argument_group("Visualization Options")
+ viz_group.add_argument(
+ "--type",
+ choices=[t.value for t in VisualizationType],
+ required=True,
+ help="Type of visualization to generate",
+ )
+ viz_group.add_argument(
+ "--entity", help="Name of the entity to visualize (function, class, file, etc.)"
+ )
+ viz_group.add_argument(
+ "--max-depth",
+ type=int,
+ default=5,
+ help="Maximum depth for recursive visualizations",
+ )
+ viz_group.add_argument(
+ "--ignore-external", action="store_true", help="Ignore external dependencies"
+ )
+ viz_group.add_argument(
+ "--severity",
+ choices=[s.value for s in IssueSeverity],
+ help="Filter issues by severity",
+ )
+ viz_group.add_argument("--path-filter", help="Filter by file path")
+
+ # PR options
+ pr_group = parser.add_argument_group("PR Options")
+ pr_group.add_argument("--pr-number", type=int, help="PR number to analyze")
+ pr_group.add_argument(
+ "--base-branch", default="main", help="Base branch for comparison"
+ )
+
+ # Output options
+ output_group = parser.add_argument_group("Output Options")
+ output_group.add_argument(
+ "--output-format",
+ choices=[f.value for f in OutputFormat],
+ default="json",
+ help="Output format for the visualization",
+ )
+ output_group.add_argument(
+ "--output-directory", help="Directory to save visualizations"
+ )
+ output_group.add_argument(
+ "--layout",
+ choices=["spring", "kamada_kawai", "spectral"],
+ default="spring",
+ help="Layout algorithm for graph visualization",
+ )
+
+ args = parser.parse_args()
+
+ # Create visualizer configuration
+ config = VisualizationConfig(
+ max_depth=args.max_depth,
+ ignore_external=args.ignore_external,
+ output_format=OutputFormat(args.output_format),
+ output_directory=args.output_directory,
+ layout_algorithm=args.layout,
+ )
+
+ # Create codebase analyzer if needed for PR comparison
+ analyzer = None
+ if args.type == VisualizationType.PR_COMPARISON.value or args.pr_number:
+ analyzer = CodebaseAnalyzer(
+ repo_url=args.repo_url,
+ repo_path=args.repo_path,
+ base_branch=args.base_branch,
+ pr_number=args.pr_number,
+ language=args.language,
+ )
+
+ # Create visualizer
+ visualizer = CodebaseVisualizer(analyzer=analyzer, config=config)
+
+ # Generate visualization based on type
+ viz_type = VisualizationType(args.type)
+ result = None
+
+ if viz_type == VisualizationType.CALL_GRAPH:
+ if not args.entity:
+ logger.error("Entity name required for call graph visualization")
+ sys.exit(1)
+
+ result = visualizer.visualize_call_graph(args.entity)
+
+ elif viz_type == VisualizationType.DEPENDENCY_GRAPH:
+ if not args.entity:
+ logger.error("Entity name required for dependency graph visualization")
+ sys.exit(1)
+
+ result = visualizer.visualize_dependency_graph(args.entity)
+
+ elif viz_type == VisualizationType.BLAST_RADIUS:
+ if not args.entity:
+ logger.error("Entity name required for blast radius visualization")
+ sys.exit(1)
+
+ result = visualizer.visualize_blast_radius(args.entity)
+
+ elif viz_type == VisualizationType.CLASS_METHODS:
+ if not args.entity:
+ logger.error("Class name required for class methods visualization")
+ sys.exit(1)
+
+ result = visualizer.visualize_class_methods(args.entity)
+
+ elif viz_type == VisualizationType.MODULE_DEPENDENCIES:
+ if not args.entity:
+ logger.error("Module path required for module dependencies visualization")
+ sys.exit(1)
+
+ result = visualizer.visualize_module_dependencies(args.entity)
+
+ elif viz_type == VisualizationType.DEAD_CODE:
+ result = visualizer.visualize_dead_code(args.path_filter)
+
+ elif viz_type == VisualizationType.CYCLOMATIC_COMPLEXITY:
+ result = visualizer.visualize_cyclomatic_complexity(args.path_filter)
+
+ elif viz_type == VisualizationType.ISSUES_HEATMAP:
+ severity = IssueSeverity(args.severity) if args.severity else None
+ result = visualizer.visualize_issues_heatmap(severity, args.path_filter)
+
+ elif viz_type == VisualizationType.PR_COMPARISON:
+ if not args.pr_number:
+ logger.error("PR number required for PR comparison visualization")
+ sys.exit(1)
+
+ result = visualizer.visualize_pr_comparison()
+
+ # Output result
+ if result:
+ logger.info(f"Visualization completed: {result}")
+ else:
+ logger.error("Failed to generate visualization")
+ sys.exit(1)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/__init__.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/__init__.py
new file mode 100644
index 000000000..5b9d135f7
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/__init__.py
@@ -0,0 +1,6 @@
+"""
+Dependency Graph Visualization Module
+
+This module provides tools for visualizing dependency relationships and impact analysis in a codebase.
+"""
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/blast_radius.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/blast_radius.py
new file mode 100644
index 000000000..42b039632
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/blast_radius.py
@@ -0,0 +1,119 @@
+import codegen
+import networkx as nx
+from codegen import Codebase
+from codegen.sdk.core.dataclasses.usage import Usage
+from codegen.sdk.core.function import PyFunction
+from codegen.sdk.core.symbol import PySymbol
+
+# Create a directed graph for visualizing relationships between code elements
+G = nx.DiGraph()
+
+# Maximum depth to traverse in the call graph to prevent infinite recursion
+MAX_DEPTH = 5
+
+# Define colors for different types of nodes in the visualization
+COLOR_PALETTE = {
+ "StartFunction": "#9cdcfe", # Starting function (light blue)
+ "PyFunction": "#a277ff", # Python functions (purple)
+ "PyClass": "#ffca85", # Python classes (orange)
+ "ExternalModule": "#f694ff", # External module imports (pink)
+ "HTTP_METHOD": "#ffca85", # HTTP method handlers (orange)
+}
+
+# List of common HTTP method names to identify route handlers
+HTTP_METHODS = ["get", "put", "patch", "post", "head", "delete"]
+
+
+def generate_edge_meta(usage: Usage) -> dict:
+ """
+ Generate metadata for graph edges based on a usage relationship.
+
+ Args:
+ usage: A Usage object representing how a symbol is used
+
+ Returns:
+ dict: Edge metadata including source location and symbol info
+ """
+ return {"name": usage.match.source, "file_path": usage.match.filepath, "start_point": usage.match.start_point, "end_point": usage.match.end_point, "symbol_name": usage.match.__class__.__name__}
+
+
+def is_http_method(symbol: PySymbol) -> bool:
+ """
+ Check if a symbol represents an HTTP method handler.
+
+ Args:
+ symbol: A Python symbol to check
+
+ Returns:
+ bool: True if symbol is an HTTP method handler
+ """
+ if isinstance(symbol, PyFunction) and symbol.is_method:
+ return symbol.name in HTTP_METHODS
+ return False
+
+
+def create_blast_radius_visualization(symbol: PySymbol, depth: int = 0):
+ """
+ Recursively build a graph visualization showing how a symbol is used.
+ Shows the "blast radius" - everything that would be affected by changes.
+
+ Args:
+ symbol: Starting symbol to analyze
+ depth: Current recursion depth
+ """
+ # Stop recursion if we hit max depth
+ if depth >= MAX_DEPTH:
+ return
+
+ # Process each usage of the symbol
+ for usage in symbol.usages:
+ usage_symbol = usage.usage_symbol
+
+ # Determine node color based on symbol type
+ if is_http_method(usage_symbol):
+ color = COLOR_PALETTE.get("HTTP_METHOD")
+ else:
+ color = COLOR_PALETTE.get(usage_symbol.__class__.__name__, "#f694ff")
+
+ # Add node and edge to graph
+ G.add_node(usage_symbol, color=color)
+ G.add_edge(symbol, usage_symbol, **generate_edge_meta(usage))
+
+ # Recurse to process usages of this symbol
+ create_blast_radius_visualization(usage_symbol, depth + 1)
+
+
+@codegen.function("visualize-function-blast-radius")
+def run(codebase: Codebase):
+ """
+ Generate a visualization showing the blast radius of changes to a function.
+
+ This codemod:
+ 1. Identifies all usages of a target function
+ 2. Creates a graph showing how the function is used throughout the codebase
+ 3. Highlights HTTP method handlers and different types of code elements
+ """
+ global G
+ G = nx.DiGraph()
+
+ # Get the target function to analyze
+ target_func = codebase.get_function("export_asset")
+
+ # Add starting function to graph with special color
+ G.add_node(target_func, color=COLOR_PALETTE.get("StartFunction"))
+
+ # Build the visualization starting from target function
+ create_blast_radius_visualization(target_func)
+
+ print(G)
+ print("Use codegen.sh to visualize the graph!")
+
+
+if __name__ == "__main__":
+ print("Initializing codebase...")
+ codebase = Codebase.from_repo("codegen-oss/posthog", commit="b174f2221ea4ae50e715eb6a7e70e9a2b0760800", language="python")
+ print(f"Codebase with {len(codebase.files)} files and {len(codebase.functions)} functions.")
+ print("Creating graph...")
+
+ run(codebase)
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/dependency_trace.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/dependency_trace.py
new file mode 100644
index 000000000..85448ac4f
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/dependency_trace.py
@@ -0,0 +1,83 @@
+import codegen
+import networkx as nx
+from codegen import Codebase
+from codegen.sdk.core.class_definition import Class
+from codegen.sdk.core.import_resolution import Import
+from codegen.sdk.core.symbol import Symbol
+
+G = nx.DiGraph()
+
+IGNORE_EXTERNAL_MODULE_CALLS = True
+IGNORE_CLASS_CALLS = False
+MAX_DEPTH = 10
+
+COLOR_PALETTE = {
+ "StartFunction": "#9cdcfe", # Light blue for the starting function
+ "PyFunction": "#a277ff", # Purple for Python functions
+ "PyClass": "#ffca85", # Orange for Python classes
+ "ExternalModule": "#f694ff", # Pink for external module references
+}
+
+# Dictionary to track visited nodes and prevent cycles
+visited = {}
+
+
+def create_dependencies_visualization(symbol: Symbol, depth: int = 0):
+ """Creates a visualization of symbol dependencies in the codebase
+
+ Recursively traverses the dependency tree of a symbol (function, class, etc.)
+ and creates a directed graph representation. Dependencies can be either direct
+ symbol references or imports.
+
+ Args:
+ symbol (Symbol): The starting symbol whose dependencies will be mapped
+ depth (int): Current depth in the recursive traversal
+ """
+ if depth >= MAX_DEPTH:
+ return
+
+ for dep in symbol.dependencies:
+ dep_symbol = None
+
+ if isinstance(dep, Symbol):
+ dep_symbol = dep
+ elif isinstance(dep, Import):
+ dep_symbol = dep.resolved_symbol if dep.resolved_symbol else None
+
+ if dep_symbol:
+ G.add_node(dep_symbol, color=COLOR_PALETTE.get(dep_symbol.__class__.__name__, "#f694ff"))
+ G.add_edge(symbol, dep_symbol)
+
+ if not isinstance(dep_symbol, Class):
+ create_dependencies_visualization(dep_symbol, depth + 1)
+
+
+@codegen.function("visualize-symbol-dependencies")
+def run(codebase: Codebase):
+ """Generate a visualization of symbol dependencies in a codebase.
+
+ This codemod:
+ 1. Creates a directed graph of symbol dependencies starting from a target function
+ 2. Tracks relationships between functions, classes, and imports
+ 3. Generates a visual representation of the dependency hierarchy
+ """
+ global G
+ G = nx.DiGraph()
+
+ target_func = codebase.get_function("get_query_runner")
+ G.add_node(target_func, color=COLOR_PALETTE.get("StartFunction"))
+
+ create_dependencies_visualization(target_func)
+
+ print(G)
+ print("Use codegen.sh to visualize the graph!")
+
+
+if __name__ == "__main__":
+ print("Initializing codebase...")
+ codebase = Codebase.from_repo("codegen-oss/posthog", commit="b174f2221ea4ae50e715eb6a7e70e9a2b0760800", language="python")
+ print(f"Codebase with {len(codebase.files)} files and {len(codebase.functions)} functions.")
+ print("Creating graph...")
+
+ run(codebase)
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/viz_dead_code.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/viz_dead_code.py
new file mode 100644
index 000000000..17e72a5a6
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/dependency_graph/viz_dead_code.py
@@ -0,0 +1,154 @@
+from abc import ABC
+
+import networkx as nx
+
+from codegen.sdk.core.codebase import CodebaseType
+from codegen.sdk.core.function import Function
+from codegen.sdk.core.import_resolution import Import
+from codegen.sdk.core.symbol import Symbol
+from codegen.shared.enums.programming_language import ProgrammingLanguage
+from tests.shared.skills.decorators import skill, skill_impl
+from tests.shared.skills.skill import Skill
+from tests.shared.skills.skill_test import SkillTestCase, SkillTestCasePyFile
+
+PyDeadCodeTest = SkillTestCase(
+ [
+ SkillTestCasePyFile(
+ input="""
+# Live code
+def used_function():
+ return "I'm used!"
+
+class UsedClass:
+ def used_method(self):
+ return "I'm a used method!"
+
+# Dead code
+def unused_function():
+ return "I'm never called!"
+
+class UnusedClass:
+ def unused_method(self):
+ return "I'm never used!"
+
+# Second-order dead code
+def second_order_dead():
+ unused_function()
+ UnusedClass().unused_method()
+
+# More live code
+def another_used_function():
+ return used_function()
+
+# Main execution
+def main():
+ print(used_function())
+ print(UsedClass().used_method())
+ print(another_used_function())
+
+if __name__ == "__main__":
+ main()
+""",
+ filepath="example.py",
+ ),
+ SkillTestCasePyFile(
+ input="""
+# This file should be ignored by the DeadCode skill
+
+from example import used_function, UsedClass
+
+def test_used_function():
+ assert used_function() == "I'm used!"
+
+def test_used_class():
+ assert UsedClass().used_method() == "I'm a used method!"
+""",
+ filepath="test_example.py",
+ ),
+ SkillTestCasePyFile(
+ input="""
+# This file contains a decorated function that should be ignored
+
+from functools import lru_cache
+
+@lru_cache
+def cached_function():
+ return "I'm cached!"
+
+# This function is dead code but should be ignored due to decoration
+@deprecated
+def old_function():
+ return "I'm old but decorated!"
+
+# This function is dead code and should be detected
+def real_dead_code():
+ return "I'm really dead!"
+""",
+ filepath="decorated_functions.py",
+ ),
+ ],
+ graph=True,
+)
+
+
+@skill(
+ eval_skill=False,
+ prompt="Show me a visualization of the call graph from my_class and filter out test files and include only the methods that have the name post, get, patch, delete",
+ uid="ec5e98c9-b57f-43f8-8b3c-af1b30bb91e6",
+)
+class DeadCode(Skill, ABC):
+ """This skill shows a visualization of the dead code in the codebase.
+ It iterates through all functions in the codebase, identifying those
+ that have no usages and are not in test files or decorated. These functions
+ are considered 'dead code' and are added to a directed graph. The skill
+ then explores the dependencies of these dead code functions, adding them to
+ the graph as well. This process helps to identify not only directly unused code
+ but also code that might only be used by other dead code (second-order dead code).
+ The resulting visualization provides a clear picture of potentially removable code,
+ helping developers to clean up and optimize their codebase.
+ """
+
+ @staticmethod
+ @skill_impl(test_cases=[PyDeadCodeTest], language=ProgrammingLanguage.PYTHON)
+ @skill_impl(test_cases=[], skip_test=True, language=ProgrammingLanguage.TYPESCRIPT)
+ def skill_func(codebase: CodebaseType):
+ # Create a directed graph to visualize dead and second-order dead code
+ G = nx.DiGraph()
+
+ # First, identify all dead code
+ dead_code: list[Function] = []
+
+ # Iterate through all functions in the codebase
+ for function in codebase.functions:
+ # Filter down functions
+ if "test" in function.file.filepath:
+ continue
+
+ if function.decorators:
+ continue
+
+ # Check if the function has no usages
+ if not function.symbol_usages:
+ # Add the function to the dead code list
+ dead_code.append(function)
+ # Add the function to the graph as dead code
+ G.add_node(function, color="red")
+
+ # # Now, find second-order dead code
+ for symbol in dead_code:
+ # Get all usages of the dead code symbol
+ for dep in symbol.dependencies:
+ if isinstance(dep, Import):
+ dep = dep.imported_symbol
+ if isinstance(dep, Symbol):
+ if "test" not in dep.name:
+ G.add_node(dep)
+ G.add_edge(symbol, dep, color="red")
+ for usage_symbol in dep.symbol_usages:
+ if isinstance(usage_symbol, Function):
+ if "test" not in usage_symbol.name:
+ G.add_edge(usage_symbol, dep)
+
+ # Visualize the graph to show dead and second-order dead code
+ codebase.visualize(G)
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/structure_graph/__init__.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/structure_graph/__init__.py
new file mode 100644
index 000000000..82dfcb765
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/structure_graph/__init__.py
@@ -0,0 +1,6 @@
+"""
+Structure Graph Visualization Module
+
+This module provides tools for visualizing code structure, directory trees, and database relationships.
+"""
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/structure_graph/graph_viz_dir_tree.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/structure_graph/graph_viz_dir_tree.py
new file mode 100644
index 000000000..67fe5e0a7
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/structure_graph/graph_viz_dir_tree.py
@@ -0,0 +1,111 @@
+from abc import ABC
+
+import networkx as nx
+
+from codegen.sdk.core.codebase import CodebaseType
+from codegen.shared.enums.programming_language import ProgrammingLanguage
+from tests.shared.skills.decorators import skill, skill_impl
+from tests.shared.skills.skill import Skill
+from tests.shared.skills.skill_test import SkillTestCase, SkillTestCasePyFile
+
+PyRepoDirTreeTest = SkillTestCase(
+ [
+ SkillTestCasePyFile(input="# Root level file", filepath="README.md"),
+ SkillTestCasePyFile(input="# Configuration file", filepath="config.yaml"),
+ SkillTestCasePyFile(
+ input="""
+def main():
+ print("Hello, World!")
+
+if __name__ == "__main__":
+ main()
+""",
+ filepath="src/main.py",
+ ),
+ SkillTestCasePyFile(
+ input="""
+class User:
+ def __init__(self, name):
+ self.name = name
+""",
+ filepath="src/models/user.py",
+ ),
+ SkillTestCasePyFile(
+ input="""
+from src.models.user import User
+
+def create_user(name):
+ return User(name)
+""",
+ filepath="src/services/user_service.py",
+ ),
+ SkillTestCasePyFile(
+ input="""
+import unittest
+from src.models.user import User
+
+class TestUser(unittest.TestCase):
+ def test_user_creation(self):
+ user = User("Alice")
+ self.assertEqual(user.name, "Alice")
+""",
+ filepath="tests/test_user.py",
+ ),
+ SkillTestCasePyFile(
+ input="""
+{
+ "name": "my-project",
+ "version": "1.0.0",
+ "description": "A sample project"
+}
+""",
+ filepath="package.json",
+ ),
+ SkillTestCasePyFile(
+ input="""
+node_modules/
+*.log
+.DS_Store
+""",
+ filepath=".gitignore",
+ ),
+ ],
+ graph=True,
+)
+
+
+@skill(eval_skill=False, prompt="Show me the directory structure of this codebase", uid="ef9a5a54-d793-4749-992d-63ea3958056b")
+class RepoDirTree(Skill, ABC):
+ """This skill displays the directory or repository tree structure of a codebase. It analyzes the file paths within the codebase and constructs a hierarchical
+ representation of the directory structure. The skill creates a visual graph where each node represents a directory or file, and edges represent the parent-child
+ relationships between directories. This visualization helps developers understand the overall organization of the codebase, making it easier to navigate and
+ manage large projects. Additionally, it can be useful for identifying potential structural issues or inconsistencies in the project layout.
+ """
+
+ @staticmethod
+ @skill_impl(test_cases=[PyRepoDirTreeTest], language=ProgrammingLanguage.PYTHON)
+ @skill_impl(test_cases=[], skip_test=True, language=ProgrammingLanguage.TYPESCRIPT)
+ def skill_func(codebase: CodebaseType):
+ # Create a directed graph
+ G = nx.DiGraph()
+
+ # Iterate over all files in the codebase
+ for file in codebase.files:
+ # Get the full filepath
+ filepath = file.filepath
+ # Split the filepath into parts
+ parts = filepath.split("/")
+
+ # Add nodes and edges to the graph
+ for i in range(len(parts)):
+ # Create a path from the root to the current part
+ path = "/".join(parts[: i + 1])
+ # Add the node for the current directory
+ G.add_node(path)
+ # If it's not the root, add an edge from the parent directory to the current directory
+ if i > 0:
+ parent_path = "/".join(parts[:i])
+ G.add_edge(parent_path, path)
+
+ codebase.visualize(G)
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/structure_graph/graph_viz_foreign_key.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/structure_graph/graph_viz_foreign_key.py
new file mode 100644
index 000000000..1f453223b
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/structure_graph/graph_viz_foreign_key.py
@@ -0,0 +1,178 @@
+from abc import ABC
+
+import networkx as nx
+
+from codegen.sdk.core.codebase import CodebaseType
+from codegen.shared.enums.programming_language import ProgrammingLanguage
+from tests.shared.skills.decorators import skill, skill_impl
+from tests.shared.skills.skill import Skill
+from tests.shared.skills.skill_test import SkillTestCase, SkillTestCasePyFile
+
+PyForeignKeyGraphTest = SkillTestCase(
+ [
+ SkillTestCasePyFile(
+ input="""
+from sqlalchemy import Column, Integer, String, ForeignKey, BigInteger
+from app.models.base import BaseModel
+
+class UserModel(BaseModel):
+ __tablename__ = 'users'
+
+ id = Column(Integer, primary_key=True)
+ name = Column(String(100), nullable=False)
+ email = Column(String(100), unique=True, nullable=False)
+
+class TaskModel(BaseModel):
+ __tablename__ = 'tasks'
+
+ id = Column(Integer, primary_key=True)
+ title = Column(String(200), nullable=False)
+ description = Column(String(500))
+ user_id = Column(Integer, ForeignKey("users.id", ondelete="CASCADE"), nullable=False)
+
+class CommentModel(BaseModel):
+ __tablename__ = 'comments'
+
+ id = Column(Integer, primary_key=True)
+ content = Column(String(500), nullable=False)
+ task_id = Column(Integer, ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False)
+ user_id = Column(Integer, ForeignKey("users.id", ondelete="CASCADE"), nullable=False)
+
+class ProjectModel(BaseModel):
+ __tablename__ = 'projects'
+
+ id = Column(Integer, primary_key=True)
+ name = Column(String(200), nullable=False)
+ description = Column(String(500))
+
+class TaskProjectModel(BaseModel):
+ __tablename__ = 'task_projects'
+
+ id = Column(Integer, primary_key=True)
+ task_id = Column(Integer, ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False)
+ project_id = Column(Integer, ForeignKey("projects.id", ondelete="CASCADE"), nullable=False)
+
+class AgentRunModel(BaseModel):
+ __tablename__ = 'agent_runs'
+
+ id = Column(BigInteger, primary_key=True)
+ task_id = Column(BigInteger, ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False)
+ agent_id = Column(BigInteger, ForeignKey("agents.id", ondelete="CASCADE"), nullable=False)
+
+class AgentModel(BaseModel):
+ __tablename__ = 'agents'
+
+ id = Column(BigInteger, primary_key=True)
+ name = Column(String(100), nullable=False)
+""",
+ filepath="app/models/schema.py",
+ )
+ ],
+ graph=True,
+)
+
+
+@skill(
+ eval_skill=False,
+ prompt="Help me analyze my data schema. I have a bunch of SQLAlchemy models with foreign keys to each other, all of them are classes like this that inherit BaseModel, like the one in this file.",
+ uid="2a5d8f4d-5f02-445e-9d00-77bdb9a0d268",
+)
+class ForeignKeyGraph(Skill, ABC):
+ """This skill helps analyze a data schema by creating a graph representation of SQLAlchemy models and their foreign key relationships.
+
+ It processes a collection of SQLAlchemy models with foreign keys referencing each other. All of these models are classes that inherit from BaseModel, similar to the one in this file. Foreign keys
+ are typically defined in the following format:
+ agent_run_id = Column(BigInteger, ForeignKey("AgentRun.id", ondelete="CASCADE"), nullable=False)
+
+ The skill iterates through all classes in the codebase, identifying those that are subclasses of BaseModel. For each relevant class, it examines the attributes to find ForeignKey definitions. It
+ then builds a mapping of these relationships.
+
+ Using this mapping, the skill constructs a directed graph where:
+ - Nodes represent the models (with the 'Model' suffix stripped from their names)
+ - Edges represent the foreign key relationships between models
+
+ This graph visualization allows for easy analysis of the data schema, showing how different models are interconnected through their foreign key relationships. The resulting graph can be used to
+ understand data dependencies, optimize queries, or refactor the database schema.
+ """
+
+ @staticmethod
+ @skill_impl(test_cases=[PyForeignKeyGraphTest], language=ProgrammingLanguage.PYTHON)
+ def skill_func(codebase: CodebaseType):
+ # Create a mapping dictionary to hold relationships
+ foreign_key_mapping = {}
+
+ # Iterate through all classes in the codebase
+ for cls in codebase.classes:
+ # Check if the class is a subclass of BaseModel and defined in the correct file
+ if cls.is_subclass_of("BaseModel") and "from app.models.base import BaseModel" in cls.file.content:
+ # Initialize an empty list for the current class
+ foreign_key_mapping[cls.name] = []
+
+ # Iterate through the attributes of the class
+ for attr in cls.attributes:
+ # Check if the attribute's source contains a ForeignKey definition
+ if "ForeignKey" in attr.source:
+ # Extract the table name from the ForeignKey string
+ start_index = attr.source.find('("') + 2
+ end_index = attr.source.find(".id", start_index)
+ if end_index != -1:
+ target_table = attr.source[start_index:end_index]
+ # Append the target table to the mapping, avoiding duplicates
+ if target_table not in foreign_key_mapping[cls.name]:
+ foreign_key_mapping[cls.name].append(target_table)
+
+ # Now foreign_key_mapping contains the desired relationships
+ # print(foreign_key_mapping)
+
+ # Create a directed graph
+ G = nx.DiGraph()
+
+ # Iterate through the foreign_key_mapping to add nodes and edges
+ for model, targets in foreign_key_mapping.items():
+ # Add the model node (strip 'Model' suffix)
+ model_name = model.replace("Model", "")
+ G.add_node(model_name)
+
+ # Add edges to the target tables
+ for target in targets:
+ G.add_node(target) # Ensure the target is also a node
+ G.add_edge(model_name, target)
+
+ # Now G contains the directed graph of models and their foreign key relationships
+ # You can visualize or analyze the graph as needed
+ codebase.visualize(G)
+
+ ##############################################################################################################
+ # IN DEGREE
+ ##############################################################################################################
+
+ # Calculate in-degrees for each node
+ in_degrees = G.in_degree()
+
+ # Create a list of nodes with their in-degree counts
+ in_degree_list = [(node, degree) for node, degree in in_degrees]
+
+ # Sort the list by in-degree in descending order
+ sorted_in_degrees = sorted(in_degree_list, key=lambda x: x[1], reverse=True)
+
+ # Print the nodes with their in-degrees
+ for node, degree in sorted_in_degrees:
+ print(f"Node: {node}, In-Degree: {degree}")
+ if degree == 0:
+ G.nodes[node]["color"] = "red"
+
+ ##############################################################################################################
+ # FIND MODELS MAPPING TO TASK
+ ##############################################################################################################
+
+ # Collect models that map to the Task model
+ models_mapping_to_task = []
+ for model, targets in foreign_key_mapping.items():
+ if "Task" in targets:
+ models_mapping_to_task.append(model)
+
+ # Print the models that map to Task
+ print("Models mapping to 'Task':")
+ for model in models_mapping_to_task:
+ print(f"> {model}")
+
diff --git a/codegen-on-oss/codegen_on_oss/analyzers/visualization/visualizer.py b/codegen-on-oss/codegen_on_oss/analyzers/visualization/visualizer.py
new file mode 100644
index 000000000..81f4f61be
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/analyzers/visualization/visualizer.py
@@ -0,0 +1,360 @@
+#!/usr/bin/env python3
+"""
+Core Visualization Module
+
+This module provides the base visualization capabilities for codebases and PR analyses.
+It defines the core classes and interfaces for generating visual representations
+of code structure, dependencies, and issues.
+"""
+
+import json
+import logging
+import os
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+try:
+ import matplotlib.pyplot as plt
+ import networkx as nx
+ from matplotlib.colors import LinearSegmentedColormap
+except ImportError:
+ logging.warning(
+ "Visualization dependencies not found. Please install them with: pip install networkx matplotlib"
+ )
+
+
+class VisualizationType(str, Enum):
+ """Types of visualizations supported by this module."""
+
+ CALL_GRAPH = "call_graph"
+ DEPENDENCY_GRAPH = "dependency_graph"
+ BLAST_RADIUS = "blast_radius"
+ CLASS_METHODS = "class_methods"
+ MODULE_DEPENDENCIES = "module_dependencies"
+ DEAD_CODE = "dead_code"
+ CYCLOMATIC_COMPLEXITY = "cyclomatic_complexity"
+ ISSUES_HEATMAP = "issues_heatmap"
+ PR_COMPARISON = "pr_comparison"
+
+
+class OutputFormat(str, Enum):
+ """Output formats for visualizations."""
+
+ JSON = "json"
+ PNG = "png"
+ SVG = "svg"
+ HTML = "html"
+ DOT = "dot"
+
+
+@dataclass
+class VisualizationConfig:
+ """Configuration for visualization generation."""
+
+ max_depth: int = 5
+ ignore_external: bool = True
+ ignore_tests: bool = True
+ node_size_base: int = 300
+ edge_width_base: float = 1.0
+ filename_filter: list[str] | None = None
+ symbol_filter: list[str] | None = None
+ output_format: OutputFormat = OutputFormat.JSON
+ output_directory: str | None = None
+ layout_algorithm: str = "spring"
+ highlight_nodes: list[str] = field(default_factory=list)
+ highlight_color: str = "#ff5555"
+ color_palette: dict[str, str] = field(
+ default_factory=lambda: {
+ "Function": "#a277ff", # Purple
+ "Class": "#ffca85", # Orange
+ "File": "#80CBC4", # Teal
+ "Module": "#81D4FA", # Light Blue
+ "Variable": "#B39DDB", # Light Purple
+ "Root": "#ef5350", # Red
+ "Warning": "#FFCA28", # Amber
+ "Error": "#EF5350", # Red
+ "Dead": "#78909C", # Gray
+ "External": "#B0BEC5", # Light Gray
+ }
+ )
+
+
+class BaseVisualizer:
+ """
+ Base visualizer providing common functionality for different visualization types.
+
+ This class implements the core operations needed for visualization, including
+ graph creation, node and edge management, and output generation.
+ """
+
+ def __init__(self, config: VisualizationConfig | None = None):
+ """
+ Initialize the BaseVisualizer.
+
+ Args:
+ config: Visualization configuration options
+ """
+ self.config = config or VisualizationConfig()
+
+ # Create visualization directory if specified
+ if self.config.output_directory:
+ os.makedirs(self.config.output_directory, exist_ok=True)
+
+ # Initialize graph for visualization
+ self.graph = nx.DiGraph()
+
+ # Tracking current visualization
+ self.current_visualization_type = None
+ self.current_entity_name = None
+
+ def _initialize_graph(self):
+ """Initialize a fresh graph for visualization."""
+ self.graph = nx.DiGraph()
+
+ def _add_node(self, node: Any, **attrs):
+ """
+ Add a node to the visualization graph with attributes.
+
+ Args:
+ node: Node object to add
+ **attrs: Node attributes
+ """
+ # Skip if node already exists
+ if self.graph.has_node(node):
+ return
+
+ # Generate node ID (memory address for unique identification)
+ node_id = id(node)
+
+ # Get node name
+ if "name" in attrs:
+ node_name = attrs["name"]
+ elif hasattr(node, "name"):
+ node_name = node.name
+ elif hasattr(node, "path"):
+ node_name = str(node.path).split("/")[-1]
+ else:
+ node_name = str(node)
+
+ # Determine node type and color
+ node_type = node.__class__.__name__
+ color = attrs.get("color", self.config.color_palette.get(node_type, "#BBBBBB"))
+
+ # Add node with attributes
+ self.graph.add_node(
+ node_id,
+ original_node=node,
+ name=node_name,
+ type=node_type,
+ color=color,
+ **attrs,
+ )
+
+ return node_id
+
+ def _add_edge(self, source: Any, target: Any, **attrs):
+ """
+ Add an edge to the visualization graph with attributes.
+
+ Args:
+ source: Source node
+ target: Target node
+ **attrs: Edge attributes
+ """
+ # Get node IDs
+ source_id = id(source)
+ target_id = id(target)
+
+ # Add edge with attributes
+ self.graph.add_edge(source_id, target_id, **attrs)
+
+ def _generate_filename(
+ self, visualization_type: VisualizationType, entity_name: str
+ ):
+ """
+ Generate a filename for the visualization.
+
+ Args:
+ visualization_type: Type of visualization
+ entity_name: Name of the entity being visualized
+
+ Returns:
+ Generated filename
+ """
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+ sanitized_name = (
+ entity_name.replace("/", "_").replace("\\", "_").replace(".", "_")
+ )
+ return f"{visualization_type.value}_{sanitized_name}_{timestamp}.{self.config.output_format.value}"
+
+ def _save_visualization(
+ self, visualization_type: VisualizationType, entity_name: str, data: Any
+ ):
+ """
+ Save a visualization to file or return it.
+
+ Args:
+ visualization_type: Type of visualization
+ entity_name: Name of the entity being visualized
+ data: Visualization data to save
+
+ Returns:
+ Path to saved file or visualization data
+ """
+ self.current_visualization_type = visualization_type
+ self.current_entity_name = entity_name
+
+ filename = self._generate_filename(visualization_type, entity_name)
+
+ if self.config.output_directory:
+ filepath = os.path.join(self.config.output_directory, filename)
+ else:
+ filepath = filename
+
+ if self.config.output_format == OutputFormat.JSON:
+ with open(filepath, "w") as f:
+ json.dump(data, f, indent=2)
+ elif self.config.output_format in [OutputFormat.PNG, OutputFormat.SVG]:
+ # Save matplotlib figure
+ plt.savefig(
+ filepath, format=self.config.output_format.value, bbox_inches="tight"
+ )
+ plt.close()
+ elif self.config.output_format == OutputFormat.DOT:
+ # Save as DOT file for Graphviz
+ try:
+ from networkx.drawing.nx_agraph import write_dot
+
+ write_dot(self.graph, filepath)
+ except ImportError:
+ logging.exception(
+ "networkx.drawing.nx_agraph not available. Install pygraphviz for DOT format."
+ )
+ return None
+
+ logging.info(f"Visualization saved to {filepath}")
+ return filepath
+
+ def _convert_graph_to_json(self):
+ """
+ Convert the networkx graph to a JSON-serializable dictionary.
+
+ Returns:
+ Dictionary representation of the graph
+ """
+ nodes = []
+ for node, attrs in self.graph.nodes(data=True):
+ # Create a serializable node
+ node_data = {
+ "id": node,
+ "name": attrs.get("name", ""),
+ "type": attrs.get("type", ""),
+ "color": attrs.get("color", "#BBBBBB"),
+ }
+
+ # Add file path if available
+ if "file_path" in attrs:
+ node_data["file_path"] = attrs["file_path"]
+
+ # Add other attributes
+ for key, value in attrs.items():
+ if key not in ["name", "type", "color", "file_path", "original_node"]:
+ if (
+ isinstance(value, str | int | float | bool | list | dict)
+ or value is None
+ ):
+ node_data[key] = value
+
+ nodes.append(node_data)
+
+ edges = []
+ for source, target, attrs in self.graph.edges(data=True):
+ # Create a serializable edge
+ edge_data = {
+ "source": source,
+ "target": target,
+ }
+
+ # Add other attributes
+ for key, value in attrs.items():
+ if (
+ isinstance(value, str | int | float | bool | list | dict)
+ or value is None
+ ):
+ edge_data[key] = value
+
+ edges.append(edge_data)
+
+ return {
+ "nodes": nodes,
+ "edges": edges,
+ "metadata": {
+ "visualization_type": self.current_visualization_type,
+ "entity_name": self.current_entity_name,
+ "timestamp": datetime.now().isoformat(),
+ "node_count": len(nodes),
+ "edge_count": len(edges),
+ },
+ }
+
+ def _plot_graph(self):
+ """
+ Plot the graph using matplotlib.
+
+ Returns:
+ Matplotlib figure
+ """
+ plt.figure(figsize=(12, 10))
+
+ # Extract node positions using specified layout algorithm
+ if self.config.layout_algorithm == "spring":
+ pos = nx.spring_layout(self.graph, seed=42)
+ elif self.config.layout_algorithm == "kamada_kawai":
+ pos = nx.kamada_kawai_layout(self.graph)
+ elif self.config.layout_algorithm == "spectral":
+ pos = nx.spectral_layout(self.graph)
+ else:
+ # Default to spring layout
+ pos = nx.spring_layout(self.graph, seed=42)
+
+ # Extract node colors
+ node_colors = [
+ attrs.get("color", "#BBBBBB") for _, attrs in self.graph.nodes(data=True)
+ ]
+
+ # Extract node sizes (can be based on some metric)
+ node_sizes = [self.config.node_size_base for _ in self.graph.nodes()]
+
+ # Draw nodes
+ nx.draw_networkx_nodes(
+ self.graph, pos, node_color=node_colors, node_size=node_sizes, alpha=0.8
+ )
+
+ # Draw edges
+ nx.draw_networkx_edges(
+ self.graph,
+ pos,
+ width=self.config.edge_width_base,
+ alpha=0.6,
+ arrows=True,
+ arrowsize=10,
+ )
+
+ # Draw labels
+ nx.draw_networkx_labels(
+ self.graph,
+ pos,
+ labels={
+ node: attrs.get("name", "")
+ for node, attrs in self.graph.nodes(data=True)
+ },
+ font_size=8,
+ font_weight="bold",
+ )
+
+ plt.title(f"{self.current_visualization_type} - {self.current_entity_name}")
+ plt.axis("off")
+
+ return plt.gcf()
diff --git a/codegen-on-oss/codegen_on_oss/bucket_store.py b/codegen-on-oss/codegen_on_oss/bucket_store.py
new file mode 100644
index 000000000..f068fd691
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/bucket_store.py
@@ -0,0 +1,26 @@
+from datetime import datetime
+from importlib.metadata import version
+from typing import TYPE_CHECKING
+
+from boto3 import client
+
+if TYPE_CHECKING:
+ from types_boto3_s3 import S3Client
+
+
+class BucketStore:
+ s3_client: "S3Client"
+
+ def __init__(self, bucket_name: str):
+ self.bucket_name = bucket_name
+ self.s3_client = client("s3")
+ self.key_prefix: str = str(version("codegen"))
+
+ def upload_file(self, local_path: str, remote_path: str) -> str:
+ key = f"{self.key_prefix}/{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}/{remote_path}"
+ self.s3_client.upload_file(
+ local_path,
+ self.bucket_name,
+ key,
+ )
+ return key
diff --git a/codegen-on-oss/codegen_on_oss/cache.py b/codegen-on-oss/codegen_on_oss/cache.py
new file mode 100644
index 000000000..6f1346a98
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/cache.py
@@ -0,0 +1,5 @@
+from pathlib import Path
+
+from platformdirs import user_cache_dir
+
+cachedir = Path(user_cache_dir("codegen-on-oss", "codegen"))
diff --git a/codegen-on-oss/codegen_on_oss/cli.py b/codegen-on-oss/codegen_on_oss/cli.py
new file mode 100644
index 000000000..c1807d13e
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/cli.py
@@ -0,0 +1,128 @@
+import sys
+from pathlib import Path
+
+import click
+from loguru import logger
+
+from codegen_on_oss.cache import cachedir
+from codegen_on_oss.metrics import MetricsProfiler
+from codegen_on_oss.outputs.csv_output import CSVOutput
+from codegen_on_oss.parser import CodegenParser
+from codegen_on_oss.sources import RepoSource, all_sources
+
+logger.remove(0)
+
+
+@click.group()
+def cli():
+ pass
+
+
+@cli.command(name="run-one")
+@click.argument("url", type=str)
+@click.option(
+ "--cache-dir",
+ type=click.Path(dir_okay=True),
+ help="Cache directory",
+ default=cachedir,
+)
+@click.option(
+ "--output-path",
+ type=click.Path(dir_okay=True),
+ help="Output path",
+ default="metrics.csv",
+)
+@click.option(
+ "--commit-hash",
+ type=str,
+ help="Commit hash to parse",
+)
+@click.option(
+ "--error-output-path",
+ type=click.Path(dir_okay=True),
+ help="Error output path",
+ default=cachedir / "errors.log",
+)
+@click.option(
+ "--debug",
+ is_flag=True,
+ help="Debug mode",
+)
+def run_one(
+ url: str,
+ cache_dir: str | Path = str(cachedir),
+ output_path: str = "metrics.csv",
+ commit_hash: str | None = None,
+ error_output_path: Path = str(cachedir / "errors.log"),
+ debug: bool = False,
+):
+ """
+ Parse a repository with codegen
+ """
+ logger.add(error_output_path, level="ERROR")
+ logger.add(sys.stdout, level="DEBUG" if debug else "INFO")
+ output = CSVOutput(MetricsProfiler.fields(), output_path)
+ metrics_profiler = MetricsProfiler(output)
+
+ parser = CodegenParser(Path(cache_dir) / "repositories", metrics_profiler)
+ parser.parse(url, commit_hash)
+
+
+@cli.command()
+@click.option(
+ "--source",
+ type=click.Choice(list(all_sources.keys())),
+ default="csv",
+)
+@click.option(
+ "--output-path",
+ type=click.Path(dir_okay=True),
+ help="Output path",
+ default="metrics.csv",
+)
+@click.option(
+ "--error-output-path",
+ type=click.Path(dir_okay=True),
+ help="Error output path",
+ default="errors.log",
+)
+@click.option(
+ "--cache-dir",
+ type=click.Path(dir_okay=True),
+ help="Cache directory",
+ default=cachedir,
+)
+@click.option(
+ "--debug",
+ is_flag=True,
+ help="Debug mode",
+)
+def run(
+ source: str,
+ output_path: str,
+ error_output_path: str,
+ cache_dir: str,
+ debug: bool,
+):
+ """
+ Run codegen parsing pipeline on repositories from a given repository source.
+ """
+ logger.add(
+ error_output_path, format="{time: HH:mm:ss} {level} {message}", level="ERROR"
+ )
+ logger.add(
+ sys.stdout,
+ format="{time: HH:mm:ss} {level} {message}",
+ level="DEBUG" if debug else "INFO",
+ )
+
+ repo_source = RepoSource.from_source_type(source)
+ output = CSVOutput(MetricsProfiler.fields(), output_path)
+ metrics_profiler = MetricsProfiler(output)
+ parser = CodegenParser(Path(cache_dir) / "repositories", metrics_profiler)
+ for repo_url, commit_hash in repo_source:
+ parser.parse(repo_url, commit_hash)
+
+
+if __name__ == "__main__":
+ cli()
diff --git a/codegen-on-oss/codegen_on_oss/codebase_analyzer.py b/codegen-on-oss/codegen_on_oss/codebase_analyzer.py
new file mode 100644
index 000000000..f4de25430
--- /dev/null
+++ b/codegen-on-oss/codegen_on_oss/codebase_analyzer.py
@@ -0,0 +1,695 @@
+#!/usr/bin/env python3
+"""
+Codebase Analyzer Module
+
+This module provides a comprehensive codebase analysis system using the Codegen SDK.
+It consolidates functionality from multiple analyzers into a single, cohesive interface
+for retrieving context about a codebase, including code quality, structure, dependencies,
+and more.
+"""
+
+import argparse
+import datetime
+import json
+import logging
+import os
+import sys
+from enum import Enum
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+
+try:
+ from codegen.configs.models.codebase import CodebaseConfig
+ from codegen.configs.models.secrets import SecretsConfig
+ from codegen.sdk.core.codebase import Codebase
+ from codegen.sdk.core.file import SourceFile
+ from codegen.sdk.core.function import Function
+ from codegen.sdk.core.import_resolution import Import
+ from codegen.sdk.core.symbol import Symbol
+ from codegen.sdk.enums import EdgeType, SymbolType
+ from codegen.shared.enums.programming_language import ProgrammingLanguage
+except ImportError as err:
+ print("Codegen SDK not found.")
+ sys.exit(1)
+
+# Configure logging
+logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+ handlers=[logging.StreamHandler()],
+)
+logger = logging.getLogger(__name__)
+
+
+class IssueSeverity(str, Enum):
+ """Severity levels for issues."""
+ CRITICAL = "critical" # Must be fixed immediately, blocks functionality
+ ERROR = "error" # Must be fixed, causes errors or undefined behavior
+ WARNING = "warning" # Should be fixed, may cause problems in future
+ INFO = "info" # Informational, could be improved but not critical
+
+
+class IssueCategory(str, Enum):
+ """Categories of issues that can be detected."""
+ # Code Quality Issues
+ UNUSED_CODE = "unused_code"
+ COMPLEXITY = "complexity"
+ STYLE = "style"
+ MAINTAINABILITY = "maintainability"
+
+ # Dependency Issues
+ CIRCULAR_DEPENDENCY = "circular_dependency"
+ DEPENDENCY_MISMATCH = "dependency_mismatch"
+ IMPORT_ERROR = "import_error"
+
+ # Implementation Issues
+ PARAMETER_ERROR = "parameter_error"
+ RETURN_TYPE_ERROR = "return_type_error"
+ IMPLEMENTATION_ERROR = "implementation_error"
+
+ # Structure Issues
+ ARCHITECTURE = "architecture"
+ COHESION = "cohesion"
+ COUPLING = "coupling"
+
+
+class CodeIssue:
+ """Represents a code issue found during analysis."""
+
+ def __init__(
+ self,
+ category: IssueCategory,
+ severity: IssueSeverity,
+ message: str,
+ file_path: Optional[str] = None,
+ line_number: Optional[int] = None,
+ symbol_name: Optional[str] = None,
+ symbol_type: Optional[str] = None,
+ code_snippet: Optional[str] = None,
+ suggestion: Optional[str] = None,
+ ):
+ self.category = category
+ self.severity = severity
+ self.message = message
+ self.file_path = file_path
+ self.line_number = line_number
+ self.symbol_name = symbol_name
+ self.symbol_type = symbol_type
+ self.code_snippet = code_snippet
+ self.suggestion = suggestion
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert the issue to a dictionary."""
+ return {
+ "category": self.category,
+ "severity": self.severity,
+ "message": self.message,
+ "file_path": self.file_path,
+ "line_number": self.line_number,
+ "symbol_name": self.symbol_name,
+ "symbol_type": self.symbol_type,
+ "code_snippet": self.code_snippet,
+ "suggestion": self.suggestion,
+ }
+
+ def __str__(self) -> str:
+ """String representation of the issue."""
+ location = f"{self.file_path}:{self.line_number}" if self.file_path else "Unknown location"
+ symbol = f"{self.symbol_type} '{self.symbol_name}'" if self.symbol_name else ""
+ return f"[{self.severity.upper()}] {self.category}: {self.message} in {location} {symbol}"
+
+
+class CodebaseAnalyzer:
+ """
+ Comprehensive codebase analyzer using Codegen SDK.
+
+ This class provides methods to analyze a codebase and extract detailed information
+ about its structure, dependencies, code quality, and more.
+ """
+
+ def __init__(
+ self,
+ repo_path: Optional[str] = None,
+ language: Optional[str] = None,
+ ):
+ """
+ Initialize the CodebaseAnalyzer.
+
+ Args:
+ repo_path: Local path to the repository to analyze
+ language: Programming language of the codebase (auto-detected if not provided)
+ """
+ self.repo_path = repo_path
+ self.language = language
+ self.codebase: Optional[Codebase] = None
+ self.issues: List[CodeIssue] = []
+
+ # Initialize the codebase if path is provided
+ if repo_path:
+ self._init_codebase(repo_path, language)
+
+ def _init_codebase(self, repo_path: str, language: Optional[str] = None):
+ """Initialize codebase from a local repository path."""
+ try:
+ # Configure the codebase
+ config = CodebaseConfig(
+ debug=False,
+ allow_external=True,
+ py_resolve_syspath=True,
+ )
+
+ secrets = SecretsConfig()
+
+ # Initialize the codebase
+ logger.info(f"Initializing codebase from {repo_path}...")
+
+ prog_lang = None
+ if language:
+ prog_lang = ProgrammingLanguage(language.upper())
+
+ self.codebase = Codebase(
+ repo_path=repo_path,
+ language=prog_lang,
+ config=config,
+ secrets=secrets
+ )
+
+ logger.info(f"Successfully initialized codebase from {repo_path}")
+
+ except Exception as e:
+ logger.error(f"Error initializing codebase from path: {e}")
+ raise
+
+ def analyze(self, output_format: str = "text", output_file: Optional[str] = None) -> Dict[str, Any]:
+ """
+ Perform a comprehensive analysis of the codebase.
+
+ Args:
+ output_format: Format of the output (text, json, html)
+ output_file: Path to the output file
+
+ Returns:
+ Dict containing the analysis results
+ """
+ if not self.codebase:
+ raise ValueError("Codebase not initialized. Please initialize the codebase first.")
+
+ # Perform all analyses
+ self._analyze_code_quality()
+ self._analyze_dependencies()
+ self._analyze_structure()
+
+ # Generate report
+ results = self._generate_report()
+
+ # Output results
+ if output_format == "json":
+ self._output_json(results, output_file)
+ elif output_format == "html":
+ self._output_html(results, output_file)
+ else:
+ self._output_text(results, output_file)
+
+ return results
+
+ def _analyze_code_quality(self):
+ """Analyze code quality issues."""
+ logger.info("Analyzing code quality...")
+
+ # Find unused functions
+ self._find_unused_functions()
+
+ # Find unused imports
+ self._find_unused_imports()
+
+ # Find functions with unused parameters
+ self._find_unused_parameters()
+
+ # Find complex functions
+ self._find_complex_functions()
+
+ logger.info(f"Found {len(self.issues)} code quality issues")
+
+ def _find_unused_functions(self):
+ """Find unused functions in the codebase."""
+ if not self.codebase:
+ return
+
+ for func in self.codebase.functions:
+ if not func.usages:
+ self.issues.append(CodeIssue(
+ category=IssueCategory.UNUSED_CODE,
+ severity=IssueSeverity.WARNING,
+ message="Function is never called",
+ file_path=func.file.file_path if hasattr(func, "file") else None,
+ symbol_name=func.name,
+ symbol_type="function",
+ suggestion="Consider removing this function or documenting why it's needed"
+ ))
+
+ def _find_unused_imports(self):
+ """Find unused imports in the codebase."""
+ if not self.codebase:
+ return
+
+ for file in self.codebase.files:
+ for import_stmt in file.imports:
+ if not import_stmt.usages:
+ self.issues.append(CodeIssue(
+ category=IssueCategory.UNUSED_CODE,
+ severity=IssueSeverity.INFO,
+ message="Import is never used",
+ file_path=file.file_path,
+ symbol_name=import_stmt.source,
+ symbol_type="import",
+ suggestion="Remove this unused import"
+ ))
+
+ def _find_unused_parameters(self):
+ """Find functions with unused parameters."""
+ if not self.codebase:
+ return
+
+ for func in self.codebase.functions:
+ for param in func.parameters:
+ # Check if parameter is used in function body
+ if param.name not in [dep.name for dep in func.dependencies]:
+ self.issues.append(CodeIssue(
+ category=IssueCategory.PARAMETER_ERROR,
+ severity=IssueSeverity.WARNING,
+ message=f"Parameter '{param.name}' is never used",
+ file_path=func.file.file_path if hasattr(func, "file") else None,
+ symbol_name=func.name,
+ symbol_type="function",
+ suggestion=f"Consider removing the unused parameter '{param.name}'"
+ ))
+
+ def _find_complex_functions(self):
+ """Find overly complex functions."""
+ if not self.codebase:
+ return
+
+ for func in self.codebase.functions:
+ if hasattr(func, 'source') and func.source:
+ line_count = func.source.count('\n') + 1
+ if line_count > 50: # Threshold for complex functions
+ self.issues.append(CodeIssue(
+ category=IssueCategory.COMPLEXITY,
+ severity=IssueSeverity.INFO,
+ message=f"Function is too complex ({line_count} lines)",
+ file_path=func.file.file_path if hasattr(func, "file") else None,
+ symbol_name=func.name,
+ symbol_type="function",
+ suggestion="Consider refactoring this function into smaller, more focused functions"
+ ))
+
+ def _analyze_dependencies(self):
+ """Analyze dependency issues."""
+ logger.info("Analyzing dependencies...")
+
+ # Find circular imports
+ self._find_circular_imports()
+
+ # Find parameter mismatches
+ self._find_parameter_mismatches()
+
+ logger.info(f"Found {len(self.issues)} dependency issues")
+
+ def _find_circular_imports(self):
+ """Find circular imports in the codebase."""
+ if not self.codebase:
+ return
+
+ # Build import graph
+ import_graph: Dict[str, List[str]] = {}
+ for file in self.codebase.files:
+ import_graph[file.file_path] = []
+ for import_stmt in file.imports:
+ if hasattr(import_stmt, 'resolved_file') and import_stmt.resolved_file:
+ import_graph[file.file_path].append(import_stmt.resolved_file.file_path)
+
+ # Check for direct circular imports
+ for file_path, imports in import_graph.items():
+ for imported_file in imports:
+ if imported_file in import_graph and file_path in import_graph[imported_file]:
+ self.issues.append(CodeIssue(
+ category=IssueCategory.CIRCULAR_DEPENDENCY,
+ severity=IssueSeverity.ERROR,
+ message=f"Circular import between {file_path} and {imported_file}",
+ file_path=file_path,
+ symbol_type="import",
+ suggestion="Refactor the code to break the circular dependency"
+ ))
+
+ def _find_parameter_mismatches(self):
+ """Find parameter mismatches in function calls."""
+ if not self.codebase:
+ return
+
+ for func in self.codebase.functions:
+ for call in func.call_sites:
+ expected_params = set(p.name for p in func.parameters)
+ actual_params = set(arg.parameter for arg in call.args if arg.parameter)
+ missing = expected_params - actual_params
+ if missing and not hasattr(func, 'has_kwargs'):
+ self.issues.append(CodeIssue(
+ category=IssueCategory.PARAMETER_ERROR,
+ severity=IssueSeverity.ERROR,
+ message=f"Function call missing parameters: {', '.join(missing)}",
+ file_path=call.file.file_path if hasattr(call, "file") else None,
+ symbol_name=func.name,
+ symbol_type="function_call",
+ suggestion=f"Add the missing parameters to the function call"
+ ))
+
+ def _analyze_structure(self):
+ """Analyze structural issues."""
+ logger.info("Analyzing code structure...")
+
+ # Find large files
+ self._find_large_files()
+
+ # Find deeply nested functions
+ self._find_deeply_nested_functions()
+
+ logger.info(f"Found {len(self.issues)} structural issues")
+
+ def _find_large_files(self):
+ """Find excessively large files."""
+ if not self.codebase:
+ return
+
+ for file in self.codebase.files:
+ if hasattr(file, 'source') and file.source:
+ line_count = file.source.count('\n') + 1
+ if line_count > 500: # Threshold for large files
+ self.issues.append(CodeIssue(
+ category=IssueCategory.MAINTAINABILITY,
+ severity=IssueSeverity.WARNING,
+ message=f"File is too large ({line_count} lines)",
+ file_path=file.file_path,
+ symbol_type="file",
+ suggestion="Consider splitting this file into multiple smaller files"
+ ))
+
+ def _find_deeply_nested_functions(self):
+ """Find deeply nested functions."""
+ if not self.codebase:
+ return
+
+ for func in self.codebase.functions:
+ if hasattr(func, 'source') and func.source:
+ # Simple heuristic: count indentation levels
+ lines = func.source.split('\n')
+ max_indent = 0
+ for line in lines:
+ if line.strip() and not line.strip().startswith('#'):
+ indent = len(line) - len(line.lstrip())
+ max_indent = max(max_indent, indent)
+
+ if max_indent > 16: # Threshold for deep nesting (4 levels with 4-space indentation)
+ self.issues.append(CodeIssue(
+ category=IssueCategory.COMPLEXITY,
+ severity=IssueSeverity.WARNING,
+ message=f"Function has deep nesting (indentation: {max_indent} spaces)",
+ file_path=func.file.file_path if hasattr(func, "file") else None,
+ symbol_name=func.name,
+ symbol_type="function",
+ suggestion="Refactor to reduce nesting by extracting code into helper functions"
+ ))
+
+ def _generate_report(self) -> Dict[str, Any]:
+ """Generate a comprehensive report of the analysis results."""
+ if not self.codebase:
+ raise ValueError("Codebase not initialized. Please initialize the codebase first.")
+
+ # Gather statistics
+ stats = {
+ "files": len(self.codebase.files),
+ "functions": len(self.codebase.functions),
+ "classes": len(self.codebase.classes),
+ "imports": len(self.codebase.imports),
+ }
+
+ # Group issues by category
+ issues_by_category: Dict[str, List[Dict[str, Any]]] = {}
+ for issue in self.issues:
+ if issue.category not in issues_by_category:
+ issues_by_category[issue.category] = []
+ issues_by_category[issue.category].append(issue.to_dict())
+
+ # Group issues by severity
+ issues_by_severity: Dict[str, List[Dict[str, Any]]] = {}
+ for issue in self.issues:
+ if issue.severity not in issues_by_severity:
+ issues_by_severity[issue.severity] = []
+ issues_by_severity[issue.severity].append(issue.to_dict())
+
+ # Create report
+ report = {
+ "metadata": {
+ "repo_path": self.repo_path,
+ "language": self.language,
+ "analysis_time": datetime.datetime.now().isoformat(),
+ },
+ "statistics": stats,
+ "issues": {
+ "total": len(self.issues),
+ "by_category": issues_by_category,
+ "by_severity": issues_by_severity,
+ },
+ }
+
+ return report
+
+ def _output_json(self, results: Dict[str, Any], output_file: Optional[str] = None):
+ """Output results in JSON format."""
+ json_str = json.dumps(results, indent=2)
+
+ if output_file:
+ with open(output_file, 'w') as f:
+ f.write(json_str)
+ logger.info(f"Results saved to {output_file}")
+ else:
+ print(json_str)
+
+ def _output_html(self, results: Dict[str, Any], output_file: Optional[str] = None):
+ """Output results in HTML format."""
+ if not output_file:
+ output_file = "codebase_analysis_report.html"
+
+ # Simple HTML template
+ html = f"""
+
+
+
+ Codebase Analysis Report
+
+
+
+ Codebase Analysis Report
+
+
+
Metadata
+
Repository: {results['metadata']['repo_path']}
+
Language: {results['metadata']['language'] or 'Auto-detected'}
+
Analysis Time: {results['metadata']['analysis_time']}
+
+
Statistics
+
Files: {results['statistics']['files']}
+
Functions: {results['statistics']['functions']}
+
Classes: {results['statistics']['classes']}
+
Imports: {results['statistics']['imports']}
+
Total Issues: {results['issues']['total']}
+
+
Issues by Severity
+ """
+
+ for severity in ['critical', 'error', 'warning', 'info']:
+ if severity in results['issues']['by_severity']:
+ issues = results['issues']['by_severity'][severity]
+ html += f"""
+
{severity.title()} ({len(issues)})
+ """
+
+ if issues:
+ for issue in issues:
+ html += f"""
+
+
{issue['category']}: {issue['message']}
+
Location: {issue['file_path'] or 'Unknown'}{':' + str(issue['line_number']) if issue['line_number'] else ''}
+ {f"
Symbol: {issue['symbol_type']} '{issue['symbol_name']}'
" if issue['symbol_name'] else ''}
+ {f"
Suggestion: {issue['suggestion']}
" if issue['suggestion'] else ''}
+
+ """
+
+ html += """
+
+
Issues by Category
+ """
+
+ for category, issues in results['issues']['by_category'].items():
+ html += f"""
+
{category.replace('_', ' ').title()} ({len(issues)})
+ """
+
+ if issues:
+ html += """
+
+
+ | Severity |
+ Message |
+ Location |
+ Symbol |
+
+ """
+
+ for issue in issues:
+ html += f"""
+
+ | {issue['severity']} |
+ {issue['message']} |
+ {issue['file_path'] or 'Unknown'}{':' + str(issue['line_number']) if issue['line_number'] else ''} |
+ {f"{issue['symbol_type']} '{issue['symbol_name']}'" if issue['symbol_name'] else ''} |
+
+ """
+
+ html += """
+
+ """
+
+ html += """
+
Codebase Analysis Report
+
+
Metadata
+
Repository: {self.results["metadata"]["repo_name"]}
+
Analysis Time: {self.results["metadata"]["analysis_time"]}
+
Language: {self.results["metadata"]["language"]}
+
+
{category.replace("_", " ").title()}
+ """
+
+ for metric_name, metric_value in metrics.items():
+ html += f"""
+
+
{metric_name.replace("_", " ").title()}
+
{json.dumps(metric_value, indent=2)}
+
+ """
+
+ html += "
@@ -116,9 +119,8 @@ def test_remove_unused_exports(tmpdir) -> None:
)
}
"""
- adj_filename = "adjacent.tsx"
# language=typescript jsx
- adj_content = """
+ ADJ_CONTENT = """
import MainComponent from 'Component'
import { SharedComponent } from 'Component'
import { StateComponent } from 'utils'
@@ -127,26 +129,79 @@ def test_remove_unused_exports(tmpdir) -> None:
return (
)
}
"""
- misc_filename = "misc.tsx"
# language=typescript jsx
- misc_content = """
+ MISC_CONTENT = """
export { UnusedComponent } from 'Component'
function Helper({ props }: HelperProps) {}
export { Helper }
"""
- import_filename = "import.tsx"
# language=typescript jsx
- import_content = """
+ IMPORT_CONTENT = """
import { UnusedComponent } from 'misc'
"""
- files = {src_filename: src_content, adj_filename: adj_content, misc_filename: misc_content, import_filename: import_content}
+ # ========== [ AFTER ] ==========
+ # language=typescript jsx
+ EXPECTED_SRC_CONTENT = """
+import { SubComponent } from 'new';
+
+export default function MainComponent() {
+ const [state, setState] = useState
()
+ return ()
+}
+
+export function UnusedComponent({ props }: UnusedProps) {
+ return (
+ Unused
+ )
+}
+"""
+ # language=typescript jsx
+ EXPECTED_NEW_CONTENT = """
+export function SubComponent({ props }: SubComponentProps) {
+ return (
+
+ )
+}
+
+function HelperComponent({ props }: HelperComponentProps) {
+ return (
+
+ )
+}
+
+export function SharedComponent({ props }: SharedComponentProps) {
+ return (
+
+ )
+}
+"""
+ # language=typescript jsx
+ EXPECTED_ADJ_CONTENT = """
+import MainComponent from 'Component'
+import { SharedComponent } from 'new'
+import { StateComponent } from 'utils'
+
+function Container(props: ContainerProps) {
+ return ()
+}
+"""
+ # language=typescript jsx
+ EXPECTED_MISC_CONTENT = """
+function Helper({ props }: HelperProps) {}
+"""
+
+ files = {"Component.tsx": SRC_CONTENT, "adjacent.tsx": ADJ_CONTENT, "misc.tsx": MISC_CONTENT, "import.tsx": IMPORT_CONTENT}
with get_codebase_session(tmpdir=tmpdir, programming_language=ProgrammingLanguage.TYPESCRIPT, files=files) as codebase:
- src_file = codebase.get_file(src_filename)
- adj_file = codebase.get_file(adj_filename)
- misc_file = codebase.get_file(misc_filename)
+ src_file = codebase.get_file("Component.tsx")
+ adj_file = codebase.get_file("adjacent.tsx")
+ misc_file = codebase.get_file("misc.tsx")
new_file = codebase.create_file("new.tsx")
sub_component = src_file.get_symbol("SubComponent")
@@ -159,20 +214,7 @@ def test_remove_unused_exports(tmpdir) -> None:
src_file.remove_unused_exports()
misc_file.remove_unused_exports()
- # Verify exports in new file
- assert "export function SubComponent" in new_file.content
- assert "function HelperComponent" in new_file.content
- assert "export function HelperComponent" not in new_file.content
- assert "export function SharedComponent" in new_file.content
-
- # Verify imports updated
- assert "import { SharedComponent } from 'new'" in adj_file.content
-
- # Verify original file exports
- assert "export default function MainComponent()" in src_file.content
- assert "function UnusedComponent" in src_file.content
- assert "export function UnusedComponent" not in src_file.content
-
- # Verify misc file exports cleaned up
- assert "export { Helper }" not in misc_file.content
- assert "export { UnusedComponent } from 'Component'" not in misc_file.content
+ assert src_file.content.strip() == EXPECTED_SRC_CONTENT.strip()
+ assert new_file.content.strip() == EXPECTED_NEW_CONTENT.strip()
+ assert adj_file.content.strip() == EXPECTED_ADJ_CONTENT.strip()
+ assert misc_file.content.strip() == EXPECTED_MISC_CONTENT.strip()
diff --git a/tests/unit/codegen/sdk/typescript/namespace/test_namespace.py b/tests/unit/codegen/sdk/typescript/namespace/test_namespace.py
index aed6271b9..ab0764f76 100644
--- a/tests/unit/codegen/sdk/typescript/namespace/test_namespace.py
+++ b/tests/unit/codegen/sdk/typescript/namespace/test_namespace.py
@@ -65,8 +65,7 @@ def test_namespace_basic_symbols(tmpdir) -> None:
assert namespace.get_symbol("privateVar") is None # private not accessible
# Test symbols collection
- assert len(namespace.symbols) == 2 # only exported symbols
- assert all(symbol.is_exported for symbol in namespace.symbols)
+ assert len(namespace.symbols) == 3
def test_namespace_recursive_symbol_lookup(tmpdir) -> None:
@@ -124,44 +123,6 @@ def test_namespace_functions(tmpdir) -> None:
assert all(func.is_exported for func in namespace.functions)
-def test_namespace_function_full_name(tmpdir) -> None:
- """Test getting functions using full names."""
- FILE_NAME = "test.ts"
- # language=typescript
- FILE_CONTENT = """
- namespace Outer {
- export function shared() { return 1; }
- export namespace Inner {
- export function shared() { return 2; }
- export function unique() { return 3; }
- }
- }
- """
- with get_codebase_session(tmpdir=tmpdir, programming_language=ProgrammingLanguage.TYPESCRIPT, files={FILE_NAME: FILE_CONTENT}) as codebase:
- namespace: TSNamespace = codebase.get_symbol("Outer")
- assert namespace is not None
-
- # Test getting functions by local name
- outer_shared = namespace.get_function("shared", recursive=False)
- assert outer_shared is not None
- inner_shared = namespace.get_function("shared", recursive=True)
- assert inner_shared is not None
- # Without full names, we might get either shared function
- assert outer_shared == inner_shared
-
- # Test getting functions by full name
- outer_shared = namespace.get_function("shared", use_full_name=True)
- assert outer_shared is not None
- inner_shared = namespace.get_function("Inner.shared", use_full_name=True)
- assert inner_shared is not None
- inner_unique = namespace.get_function("Inner.unique", use_full_name=True)
- assert inner_unique is not None
-
- # Test non-existent paths
- assert namespace.get_function("NonExistent.shared", use_full_name=True) is None
- assert namespace.get_function("Inner.NonExistent", use_full_name=True) is None
-
-
def test_namespace_function_overloading(tmpdir) -> None:
"""Test function overloading within namespace."""
FILE_NAME = "test.ts"
@@ -333,3 +294,46 @@ def test_namespace_nested_deep(tmpdir) -> None:
assert len(nested) == 2 # Should find B and C
assert all(isinstance(ns, TSNamespace) for ns in nested)
assert {ns.name for ns in nested} == {"B", "C"}
+
+
+def test_namespace_imports(tmpdir) -> None:
+ """Test importing and using namespaces."""
+ FILE_NAME_1 = "math.ts"
+ # language=typescript
+ FILE_CONTENT_1 = """
+ export namespace Math {
+ export const PI = 3.14159;
+ export function square(x: number) { return x * x; }
+
+ export namespace Advanced {
+ export function cube(x: number) { return x * x * x; }
+ }
+ }
+ """
+
+ FILE_NAME_2 = "app.ts"
+ # language=typescript
+ FILE_CONTENT_2 = """
+ import { Math } from './math';
+
+ console.log(Math.PI);
+ console.log(Math.square(5));
+ console.log(Math.Advanced.cube(3));
+ """
+
+ with get_codebase_session(tmpdir=tmpdir, programming_language=ProgrammingLanguage.TYPESCRIPT, files={FILE_NAME_1: FILE_CONTENT_1, FILE_NAME_2: FILE_CONTENT_2}) as codebase:
+ math_ns = codebase.get_symbol("Math")
+ assert math_ns is not None
+ assert math_ns.name == "Math"
+
+ # Test namespace import resolution
+ file2 = codebase.get_file(FILE_NAME_2)
+ math_import = file2.get_import("Math")
+ assert math_import is not None
+ assert math_import.is_namespace_import
+
+ # Test nested namespace access
+ advanced = math_ns.get_namespace("Advanced")
+ assert advanced is not None
+ assert advanced.name == "Advanced"
+ assert advanced.get_function("cube") is not None
diff --git a/tests/unit/codegen/sdk/typescript/namespace/test_namespace_complex_examples.py b/tests/unit/codegen/sdk/typescript/namespace/test_namespace_complex_examples.py
index 3dfa77e28..9af4baf6f 100644
--- a/tests/unit/codegen/sdk/typescript/namespace/test_namespace_complex_examples.py
+++ b/tests/unit/codegen/sdk/typescript/namespace/test_namespace_complex_examples.py
@@ -133,3 +133,42 @@ def test_namespace_validators(tmpdir) -> None:
# Verify non-exported items are not accessible
assert namespace.get_symbol("lettersRegexp") is None
assert namespace.get_symbol("numberRegexp") is None
+
+
+def test_namespace_wildcard_import(tmpdir) -> None:
+ """Test wildcard imports with namespaces."""
+ FILE_NAME_1 = "utils.ts"
+ # language=typescript
+ FILE_CONTENT_1 = """
+ export namespace Utils {
+ export const helper1 = () => "help1";
+ export const helper2 = () => "help2";
+ const internal = () => "internal";
+ }
+ """
+
+ FILE_NAME_2 = "app.ts"
+ # language=typescript
+ FILE_CONTENT_2 = """
+ import * as AllUtils from './utils';
+
+ function test() {
+ console.log(AllUtils.Utils.helper1());
+ console.log(AllUtils.Utils.helper2());
+ }
+ """
+
+ with get_codebase_session(tmpdir=tmpdir, programming_language=ProgrammingLanguage.TYPESCRIPT, files={FILE_NAME_1: FILE_CONTENT_1, FILE_NAME_2: FILE_CONTENT_2}) as codebase:
+ utils_file = codebase.get_file(FILE_NAME_1)
+ app_file = codebase.get_file(FILE_NAME_2)
+
+ # Verify namespace import
+ utils_import = app_file.get_import("AllUtils")
+ assert utils_import is not None
+ assert utils_import.namespace == "AllUtils"
+
+ # Verify access to exported symbols
+ utils_ns = utils_file.get_symbol("Utils")
+ assert "helper1" in utils_ns.valid_import_names
+ assert "helper2" in utils_ns.valid_import_names
+ assert "internal" not in utils_ns.valid_import_names
diff --git a/tests/unit/codegen/sdk/typescript/namespace/test_namespace_modifications.py b/tests/unit/codegen/sdk/typescript/namespace/test_namespace_modifications.py
new file mode 100644
index 000000000..fc592beac
--- /dev/null
+++ b/tests/unit/codegen/sdk/typescript/namespace/test_namespace_modifications.py
@@ -0,0 +1,183 @@
+from typing import TYPE_CHECKING
+
+import pytest
+
+from codegen.sdk.codebase.factory.get_session import get_codebase_session
+from codegen.shared.enums.programming_language import ProgrammingLanguage
+
+if TYPE_CHECKING:
+ from codegen.sdk.typescript.namespace import TSNamespace
+
+
+def test_namespace_add_symbol(tmpdir) -> None:
+ """Test adding symbols to namespace."""
+ FILE_NAME = "test.ts"
+ # language=typescript
+ FILE_CONTENT = """
+ namespace MyNamespace {
+ export const x = 1;
+ }
+ """
+ with get_codebase_session(tmpdir=tmpdir, programming_language=ProgrammingLanguage.TYPESCRIPT, files={FILE_NAME: FILE_CONTENT}) as codebase:
+ file = codebase.get_file("test.ts")
+ namespace: TSNamespace = codebase.get_symbol("MyNamespace")
+
+ # 1. a) Add new symbol from object, then manually remove the original symbol from the file
+ # 1. b) Add new symbol by moving operation
+ file.add_symbol_from_source(source="const ya = 2")
+ codebase.ctx.commit_transactions()
+ new_const = file.get_symbol("ya")
+
+ # Store original location
+
+ # Add to namespace and remove from original location
+ namespace.add_symbol(new_const, should_export=True)
+
+ codebase.ctx.commit_transactions()
+
+ # Get fresh reference to namespace
+ namespace: TSNamespace = codebase.get_symbol("MyNamespace")
+
+ # Verify symbols were moved correctly
+ assert namespace.get_symbol("ya") is not None
+ assert namespace.get_symbol("ya").export is not None
+
+ # 2. Add new symbol from string
+ code = "const z = 3"
+ namespace.add_symbol_from_source(code)
+ codebase.ctx.commit_transactions()
+ namespace: TSNamespace = codebase.get_symbol("MyNamespace")
+
+ code_symbol = namespace.get_symbol("z", get_private=True)
+ # Verify exported symbol
+ assert code_symbol is not None
+ assert code_symbol.name == "z"
+
+ assert len(namespace.symbols) == 3
+ assert {s.name for s in namespace.symbols} == {"x", "ya", "z"}
+
+
+def test_namespace_remove_symbol(tmpdir) -> None:
+ """Test removing symbols from namespace."""
+ FILE_NAME = "test.ts"
+ # language=typescript
+ FILE_CONTENT = """
+ namespace MyNamespace {
+ export const x = 1;
+ export const y = 2;
+ }
+ """
+ with get_codebase_session(tmpdir=tmpdir, programming_language=ProgrammingLanguage.TYPESCRIPT, files={FILE_NAME: FILE_CONTENT}) as codebase:
+ namespace: TSNamespace = codebase.get_symbol("MyNamespace")
+
+ # Remove existing symbol
+ removed = namespace.remove_symbol("x")
+ codebase.ctx.commit_transactions()
+ assert removed is not None
+ assert removed.name == "x"
+
+ # Verify symbol was removed
+ assert namespace.get_symbol("x") is None
+ assert len(namespace.symbols) == 1
+ assert namespace.symbols[0].name == "y"
+
+ # Try removing non-existent symbol
+ assert namespace.remove_symbol("z") is None
+
+
+def test_namespace_rename(tmpdir) -> None:
+ """Test renaming namespace."""
+ FILE_NAME = "test.ts"
+ # language=typescript
+ FILE_CONTENT = """
+ namespace OldName {
+ export const x = 1;
+ }
+ """
+ with get_codebase_session(tmpdir=tmpdir, programming_language=ProgrammingLanguage.TYPESCRIPT, files={FILE_NAME: FILE_CONTENT}) as codebase:
+ namespace: TSNamespace = codebase.get_symbol("OldName")
+
+ # Rename namespace
+ namespace.rename("NewName")
+ codebase.ctx.commit_transactions()
+
+ # Verify rename
+ namespace: TSNamespace = codebase.get_symbol("NewName")
+ assert namespace.name == "NewName"
+ assert codebase.get_symbol("NewName") is namespace
+ assert codebase.get_symbol("OldName", optional=True) is None
+
+
+def test_namespace_export_symbol(tmpdir) -> None:
+ """Test exporting symbols in namespace."""
+ FILE_NAME = "test.ts"
+ # language=typescript
+ FILE_CONTENT = """
+ namespace ExportTest {
+ export const external = 123;
+ const internal = 123;
+ }
+ """
+ with get_codebase_session(tmpdir=tmpdir, programming_language=ProgrammingLanguage.TYPESCRIPT, files={FILE_NAME: FILE_CONTENT}) as codebase:
+ namespace: TSNamespace = codebase.get_symbol("ExportTest")
+
+ # Export internal symbol
+ namespace.export_symbol("internal")
+ codebase.ctx.commit_transactions()
+
+ # Verify export
+ namespace: TSNamespace = codebase.get_symbol("ExportTest")
+ internal = namespace.get_symbol("internal")
+ assert internal is not None
+ assert all(symbol.is_exported for symbol in namespace.symbols)
+
+ # Export already exported symbol (no change)
+ namespace.export_symbol("external")
+ codebase.ctx.commit_transactions()
+
+ namespace: TSNamespace = codebase.get_symbol("ExportTest")
+ external = namespace.get_symbol("external")
+ assert external is not None
+ assert external.is_exported
+
+
+@pytest.mark.skip("TODO: Symbol Animals is ambiguous in codebase - more than one instance")
+def test_namespace_merging(tmpdir) -> None:
+ """Test TypeScript namespace merging functionality."""
+ FILE_NAME = "test.ts"
+ # language=typescript
+ FILE_CONTENT = """
+ namespace Animals {
+ export class Dog { bark() {} }
+ }
+
+ namespace Animals { // Merge with previous namespace
+ export class Cat { meow() {} }
+ }
+
+ namespace Plants { // Different namespace, should not merge
+ export class Tree {}
+ }
+ """
+ with get_codebase_session(tmpdir=tmpdir, programming_language=ProgrammingLanguage.TYPESCRIPT, files={FILE_NAME: FILE_CONTENT}) as codebase:
+ animals = codebase.get_symbol("Animals")
+ assert animals is not None
+
+ # Test merged namespace access
+ assert animals.get_class("Dog") is not None
+ assert animals.get_class("Cat") is not None
+
+ # Verify merged namespaces
+ assert len(animals.merged_namespaces) == 1
+ merged = animals.merged_namespaces[0]
+ assert merged.name == "Animals"
+ assert merged != animals
+
+ # Verify all symbols accessible
+ all_symbols = animals.symbols
+ assert len(all_symbols) == 2
+ assert {s.name for s in all_symbols} == {"Dog", "Cat"}
+
+ # Verify non-merged namespace
+ plants = codebase.get_symbol("Plants")
+ assert len(plants.merged_namespaces) == 0
diff --git a/tests/unit/codegen/sdk/typescript/namespace/test_namespace_usage.py b/tests/unit/codegen/sdk/typescript/namespace/test_namespace_usage.py
new file mode 100644
index 000000000..9f72250b0
--- /dev/null
+++ b/tests/unit/codegen/sdk/typescript/namespace/test_namespace_usage.py
@@ -0,0 +1,103 @@
+from codegen.sdk.codebase.factory.get_session import get_codebase_session
+from codegen.sdk.core.dataclasses.usage import UsageType
+from codegen.shared.enums.programming_language import ProgrammingLanguage
+
+
+def test_namespace_same_file_usage(tmpdir) -> None:
+ """Test namespace usage within the same file."""
+ # language=typescript
+ content = """
+ namespace MathUtils {
+ export const PI = 3.14159;
+ export function square(x: number) { return x * x; }
+ }
+
+ function calculateArea(radius: number) {
+ return MathUtils.PI * MathUtils.square(radius);
+ }
+ """
+ with get_codebase_session(tmpdir=tmpdir, files={"test.ts": content}, programming_language=ProgrammingLanguage.TYPESCRIPT) as codebase:
+ file = codebase.get_file("test.ts")
+
+ namespace = file.get_symbol("MathUtils")
+ pi = namespace.get_symbol("PI")
+ square = namespace.get_symbol("square")
+ calc_area = file.get_function("calculateArea")
+
+ # Check if namespace is in valid_import_names
+ assert "MathUtils" in file.valid_symbol_names
+ assert "MathUtils" in namespace.valid_import_names
+ assert len(namespace.valid_import_names) == 3 # MathUtils, PI, and square
+
+ # Check usages
+ assert {calc_area}.issubset(namespace.symbol_usages)
+
+ # PI has direct usage (export) and chained usage (in calculateArea)
+ assert set(pi.symbol_usages(UsageType.DIRECT)) == {pi.export}
+ assert set(pi.symbol_usages(UsageType.CHAINED)) == {calc_area}
+ assert set(pi.symbol_usages) == {pi.export, calc_area}
+
+ # square has direct usage (export) and chained usage (in calculateArea)
+ assert set(square.symbol_usages(UsageType.DIRECT)) == {square.export}
+ assert set(square.symbol_usages(UsageType.CHAINED)) == {calc_area}
+ assert set(square.symbol_usages) == {square.export, calc_area}
+
+ # Verify attribute resolution
+ assert namespace.resolve_attribute("PI") == pi.export
+ assert namespace.resolve_attribute("square") == square.export
+
+
+def test_namespace_cross_file_usage(tmpdir) -> None:
+ """Test namespace usage across files with imports."""
+ # language=typescript
+ content1 = """
+ export namespace MathUtils {
+ export const PI = 3.14159;
+ export function square(x: number) { return x * x; }
+ const internal = 123; // not exported
+ }
+ """
+ # language=typescript
+ content2 = """
+ import { MathUtils } from './file1';
+
+ function calculateArea(radius: number) {
+ return MathUtils.PI * MathUtils.square(radius);
+ }
+
+ function calculateVolume(radius: number) {
+ const area = calculateArea(radius);
+ return area * radius;
+ }
+ """
+ with get_codebase_session(tmpdir=tmpdir, files={"file1.ts": content1, "file2.ts": content2}, programming_language=ProgrammingLanguage.TYPESCRIPT) as codebase:
+ file1 = codebase.get_file("file1.ts")
+ file2 = codebase.get_file("file2.ts")
+
+ # Get symbols
+ namespace = file1.get_symbol("MathUtils")
+ pi = namespace.get_symbol("PI")
+ square = namespace.get_symbol("square")
+ internal = namespace.get_symbol("internal")
+ calc_area = file2.get_function("calculateArea")
+ calc_volume = file2.get_function("calculateVolume")
+ namespace_import = file2.get_import("MathUtils")
+
+ # Check namespace visibility
+ assert "MathUtils" in namespace.valid_import_names
+ assert "PI" in namespace.valid_import_names
+ assert "square" in namespace.valid_import_names
+ assert "internal" not in namespace.valid_import_names
+ assert internal is None # private symbol not accessible
+
+ # Check direct vs chained usages
+ assert {namespace.export}.issubset(namespace.symbol_usages(UsageType.DIRECT))
+ assert {namespace.export, calc_area}.issubset(namespace.symbol_usages)
+ assert {pi.export}.issubset(pi.symbol_usages(UsageType.DIRECT))
+ assert {pi.export, calc_area}.issubset(pi.symbol_usages)
+ assert {calc_area}.issubset(square.symbol_usages(UsageType.CHAINED))
+
+ # Verify attribute resolution
+ assert namespace.resolve_attribute("PI") == pi.export
+ assert namespace.resolve_attribute("square") == square.export
+ assert namespace.resolve_attribute("internal") is None
diff --git a/tests/unit/codegen/sdk/typescript/tsx/test_tsx_edit.py b/tests/unit/codegen/sdk/typescript/tsx/test_tsx_edit.py
index 6f21af839..a7147bf3a 100644
--- a/tests/unit/codegen/sdk/typescript/tsx/test_tsx_edit.py
+++ b/tests/unit/codegen/sdk/typescript/tsx/test_tsx_edit.py
@@ -333,7 +333,7 @@ def test_tsx_move_component(tmpdir) -> None:
ctx.commit_transactions()
assert "export function FooBar" in new_file.content
- assert "export function MyFooBar" in new_file.content
+ assert "function MyFooBar" in new_file.content
assert "import { FooBar } from 'new'" in original_file.content
assert "import { MyFooBar } from 'new'" not in original_file.content
diff --git a/tests/unit/codegen/sdk/typescript/tsx/test_tsx_parsing.py b/tests/unit/codegen/sdk/typescript/tsx/test_tsx_parsing.py
index af2f32446..813102927 100644
--- a/tests/unit/codegen/sdk/typescript/tsx/test_tsx_parsing.py
+++ b/tests/unit/codegen/sdk/typescript/tsx/test_tsx_parsing.py
@@ -105,7 +105,7 @@ def test_tsx_file_type_validation(tmpdir) -> None:
test_component.move_to_file(tsx_file)
- assert "export function TestComponent" in tsx_file.content
+ assert "function TestComponent" in tsx_file.content
def test_jsx_element_attributes(tmpdir) -> None:
diff --git a/tests/unit/skills/implementations/guides/organize-your-codebase.py b/tests/unit/skills/implementations/guides/organize-your-codebase.py
index 5827d2ca5..d2e914bd2 100644
--- a/tests/unit/skills/implementations/guides/organize-your-codebase.py
+++ b/tests/unit/skills/implementations/guides/organize-your-codebase.py
@@ -416,7 +416,7 @@ def my_symbol():
SkillTestCaseTSFile(
input="",
output="""
-export function dependencyFunction() {
+function dependencyFunction() {
console.log("I'm a dependency");
}
-
-