code-compass-cli 0.1.0__py3-none-any.whl

src/visualizer/flow_tracer.py

@@ -0,0 +1,294 @@
+"""Flow tracer for visualizing code execution paths."""
+
+import os
+import re
+from pathlib import Path
+from typing import Dict, List, Set, Tuple, Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class Node:
+    """Represents a node in the execution flow graph."""
+
+    name: str
+    node_type: str
+    file: str
+    line: int
+
+
+class FlowTracer:
+    """Traces and visualizes execution flow in code."""
+
+    def __init__(self, repo_path: str = "."):
+        """Initialize the flow tracer.
+
+        Args:
+            repo_path: Path to repository to analyze
+        """
+        self.repo_path = Path(repo_path)
+        self.nodes: Dict[str, Node] = {}
+        self.edges: List[Tuple[str, str]] = []
+        self.definitions: Dict[str, Node] = {}
+        self.calls: List[Tuple[str, str, str, int]] = []  # (caller, callee, file, line)
+
+    def trace(self, entry_point: str) -> Dict[str, any]:
+        """Trace execution flow starting from an entry point.
+
+        Args:
+            entry_point: Name of the entry point function/method
+
+        Returns:
+            Dictionary representing the execution flow
+        """
+        # First, scan the codebase for definitions and calls
+        self._scan_codebase()
+
+        # Then build the call chain
+        visited: Set[str] = set()
+        flow = self._build_call_chain(entry_point, visited, depth=0)
+
+        return {
+            "entry_point": entry_point,
+            "flow": flow,
+            "total_calls": len(self.calls),
+        }
+
+    def _scan_codebase(self) -> None:
+        """Scan codebase to find function definitions and calls."""
+        for root, _, files in os.walk(self.repo_path):
+            # Skip venv and hidden directories
+            if "venv" in root or "/.git" in root:
+                continue
+
+            for file in files:
+                if file.endswith(".py"):
+                    filepath = Path(root) / file
+                    rel_path = filepath.relative_to(self.repo_path)
+
+                    try:
+                        with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
+                            content = f.read()
+                            self._extract_definitions(content, str(rel_path))
+                            self._extract_calls(content, str(rel_path))
+                    except (IOError, OSError):
+                        continue
+
+    def _extract_definitions(self, content: str, filepath: str) -> None:
+        """Extract function and class definitions.
+
+        Args:
+            content: File content
+            filepath: Path to the file
+        """
+        lines = content.split("\n")
+
+        # Find function definitions
+        for i, line in enumerate(lines, 1):
+            # Match function definitions
+            func_match = re.match(r"^\s*def\s+(\w+)\s*\(", line)
+            if func_match:
+                func_name = func_match.group(1)
+                self.definitions[func_name] = Node(
+                    name=func_name,
+                    node_type="function",
+                    file=filepath,
+                    line=i,
+                )
+
+            # Match class definitions
+            class_match = re.match(r"^\s*class\s+(\w+)\s*[\(:]", line)
+            if class_match:
+                class_name = class_match.group(1)
+                self.definitions[class_name] = Node(
+                    name=class_name,
+                    node_type="class",
+                    file=filepath,
+                    line=i,
+                )
+
+    def _extract_calls(self, content: str, filepath: str) -> None:
+        """Extract function and method calls.
+
+        Args:
+            content: File content
+            filepath: Path to the file
+        """
+        lines = content.split("\n")
+        
+        # Built-in functions and methods to skip
+        builtins = {
+            "append", "extend", "insert", "remove", "pop", "clear", "index", "count",
+            "sort", "reverse", "copy", "len", "str", "int", "float", "bool", "list",
+            "dict", "set", "tuple", "range", "enumerate", "zip", "map", "filter",
+            "open", "read", "write", "close", "strip", "split", "join", "replace",
+            "upper", "lower", "startswith", "endswith", "find", "format", "items",
+            "keys", "values", "get", "update", "exists", "mkdir", "rmdir", "walk",
+            "exit", "print", "input", "any", "all", "sum", "max", "min", "abs",
+            "round", "isinstance", "type", "hasattr", "getattr", "setattr", "callable",
+            "sorted", "reversed", "group", "match", "search", "sub", "compile",
+            "relative_to", "with", "except", "try", "assert"
+        }
+
+        for i, line in enumerate(lines, 1):
+            # Skip comments and docstrings
+            if line.strip().startswith("#") or '"""' in line or "'''" in line:
+                continue
+
+            # Find function calls - both direct and method calls
+            # Patterns: func_name(), obj.method(), module.func()
+            call_matches = re.findall(r"(\w+)\s*\(", line)
+
+            for match in call_matches:
+                # Skip built-in functions
+                if match in builtins:
+                    continue
+                    
+                # Get the context - find which function we're in
+                caller = self._find_current_function(content, i)
+                if caller:
+                    self.calls.append((caller, match, filepath, i))
+
+    def _find_current_function(self, content: str, line_num: int) -> Optional[str]:
+        """Find which function a line belongs to.
+
+        Args:
+            content: File content
+            line_num: Line number (1-indexed)
+
+        Returns:
+            Name of the function or None
+        """
+        lines = content.split("\n")
+        current_func = None
+        current_indent = float("inf")
+
+        for i in range(line_num - 1, -1, -1):
+            line = lines[i]
+            if line.strip() == "":
+                continue
+
+            indent = len(line) - len(line.lstrip())
+
+            # Looking for function definition at lower indentation
+            if indent < current_indent:
+                func_match = re.match(r"^\s*def\s+(\w+)\s*\(", line)
+                if func_match:
+                    return func_match.group(1)
+
+            if current_indent == float("inf"):
+                current_indent = indent
+
+        return None
+
+    def _build_call_chain(self, func_name: str, visited: Set[str], depth: int = 0) -> Dict[str, any]:
+        """Build the call chain for a function.
+
+        Args:
+            func_name: Name of the function
+            visited: Set of already visited functions
+            depth: Current depth in the call tree
+
+        Returns:
+            Dictionary representing the call chain
+        """
+        if func_name in visited or depth > 5:  # Prevent infinite recursion
+            return None
+
+        visited.add(func_name)
+
+        # Get the definition
+        definition = self.definitions.get(func_name)
+        if not definition:
+            return {
+                "name": func_name,
+                "type": "unknown",
+                "file": None,
+                "line": None,
+                "calls": [],
+            }
+
+        # Find all functions called by this function
+        called_functions = []
+        for caller, callee, call_file, call_line in self.calls:
+            if caller == func_name:
+                called_functions.append((callee, call_file, call_line))
+
+        # Remove duplicates and sort
+        called_functions = sorted(set(called_functions))
+
+        # Build call chain for each called function
+        children = []
+        for callee, call_file, call_line in called_functions:
+            child_flow = self._build_call_chain(callee, visited.copy(), depth + 1)
+            if child_flow:
+                child_flow["called_from"] = {
+                    "file": call_file,
+                    "line": call_line,
+                }
+                children.append(child_flow)
+
+        return {
+            "name": func_name,
+            "type": definition.node_type,
+            "file": definition.file,
+            "line": definition.line,
+            "calls": children,
+        }
+
+    def add_node(self, node: Node) -> None:
+        """Add a node to the flow graph.
+
+        Args:
+            node: Node to add
+        """
+        self.nodes[node.name] = node
+
+    def add_edge(self, source: str, target: str) -> None:
+        """Add an edge between two nodes.
+
+        Args:
+            source: Source node name
+            target: Target node name
+        """
+        self.edges.append((source, target))
+
+    def get_flow(self, entry_point: str) -> Dict[str, any]:
+        """Get the execution flow starting from an entry point.
+
+        Args:
+            entry_point: Name of the entry point node
+
+        Returns:
+            Dictionary representing the execution flow
+        """
+        visited: Set[str] = set()
+        flow = self._traverse(entry_point, visited)
+        return flow
+
+    def _traverse(self, node_name: str, visited: Set[str]) -> Dict[str, any]:
+        """Traverse the flow graph from a node.
+
+        Args:
+            node_name: Name of the node to traverse from
+            visited: Set of already visited nodes
+
+        Returns:
+            Dictionary representing the traversal
+        """
+        if node_name in visited or node_name not in self.nodes:
+            return {}
+
+        visited.add(node_name)
+        node = self.nodes[node_name]
+        children = [target for source, target in self.edges if source == node_name]
+
+        return {
+            "node": {
+                "name": node.name,
+                "type": node.node_type,
+                "file": node.file,
+                "line": node.line,
+            },
+            "children": [self._traverse(child, visited) for child in children],
+        }