diff --git a/frontend/src/components/Visualization/RepositoryGraph.tsx b/frontend/src/components/Visualization/RepositoryGraph.tsx index 705f02b..f73eda9 100644 --- a/frontend/src/components/Visualization/RepositoryGraph.tsx +++ b/frontend/src/components/Visualization/RepositoryGraph.tsx @@ -163,8 +163,8 @@ const RepositoryGraph = forwardRef( .data(links) .enter() .append('line') - .attr('stroke', '#95a5a6') - .attr('stroke-opacity', 0.4) + .attr('stroke', d => getLinkColor(d)) + .attr('stroke-opacity', 0.6) .attr('stroke-width', d => getLinkWidth(d)); // Create nodes @@ -295,8 +295,11 @@ const RepositoryGraph = forwardRef( const getLinkWidth = (link: Link) => { switch (link.type) { case 'import': - case 'call': + case 'imports_component': return 2; + case 'call': + case 'calls': + return 2.5; case 'contains': return 1; default: @@ -304,6 +307,22 @@ const RepositoryGraph = forwardRef( } }; + const getLinkColor = (link: Link) => { + switch (link.type) { + case 'import': + return '#3498db'; // Blue color for file-level imports + case 'imports_component': + return '#9b59b6'; // Purple color for component-level imports + case 'call': + case 'calls': + return '#e74c3c'; // Red color for function calls + case 'contains': + return '#95a5a6'; // Gray color for containment relationships + default: + return '#95a5a6'; // Default gray color + } + }; + const getNodeColor = (node: Node, colors: Record) => { // Directories have a different color if (node.type === 'directory') { @@ -333,20 +352,37 @@ const RepositoryGraph = forwardRef( } }); + // Get relation types used in the data + const usedRelationTypes = new Set(); + data.relationships.forEach(rel => { + usedRelationTypes.add(rel.type); + }); + + // Create legend group const legendGroup = svg.append('g').attr('transform', `translate(20, 20)`); + // Title for node types + legendGroup + .append('text') + .attr('x', 0) + .attr('y', 0) + .text('Node Types') + .style('font-size', '14px') + .style('font-weight', 'bold') + .style('fill', '#333'); + // Add directory type legendGroup .append('circle') .attr('cx', 10) - .attr('cy', 10) + .attr('cy', 20) .attr('r', 6) .attr('fill', '#7f8c8d'); legendGroup .append('text') .attr('x', 20) - .attr('y', 14) + .attr('y', 24) .text('Directory') .style('font-size', '12px') .style('fill', '#333'); @@ -358,14 +394,14 @@ const RepositoryGraph = forwardRef( legendGroup .append('circle') .attr('cx', 10 + Math.floor(index / 10) * 100) - .attr('cy', 10 + (index % 10) * 20) + .attr('cy', 20 + (index % 10) * 20) .attr('r', 6) .attr('fill', colors[ext]); legendGroup .append('text') .attr('x', 20 + Math.floor(index / 10) * 100) - .attr('y', 14 + (index % 10) * 20) + .attr('y', 24 + (index % 10) * 20) .text(`.${ext}`) .style('font-size', '12px') .style('fill', '#333'); @@ -378,17 +414,64 @@ const RepositoryGraph = forwardRef( legendGroup .append('circle') .attr('cx', 10 + Math.floor(index / 10) * 100) - .attr('cy', 10 + (index % 10) * 20) + .attr('cy', 20 + (index % 10) * 20) .attr('r', 6) .attr('fill', '#aaaaaa'); legendGroup .append('text') .attr('x', 20 + Math.floor(index / 10) * 100) - .attr('y', 14 + (index % 10) * 20) + .attr('y', 24 + (index % 10) * 20) .text('Other') .style('font-size', '12px') .style('fill', '#333'); + + // Find maximum Y position used by node type legend + const nodeTypeHeight = 24 + Math.min(index + 1, 10) * 20; + + // Title for relationship types with some spacing + legendGroup + .append('text') + .attr('x', 0) + .attr('y', nodeTypeHeight + 30) + .text('Relationship Types') + .style('font-size', '14px') + .style('font-weight', 'bold') + .style('fill', '#333'); + + // Link legend entries + const relationshipTypes = [ + { type: 'contains', color: '#95a5a6', label: 'Contains' }, + { type: 'import', color: '#3498db', label: 'File Import' }, + { type: 'imports_component', color: '#9b59b6', label: 'Component Import' }, + { type: 'call', color: '#e74c3c', label: 'Function Call' }, + ]; + + // Only show relationship types that are present in the data + const filteredRelationships = relationshipTypes.filter( + rel => usedRelationTypes.has(rel.type) || rel.type === 'imports_component' + ); + + // Add relationship legend + filteredRelationships.forEach((rel, i) => { + // Draw a small line instead of a circle for links + legendGroup + .append('line') + .attr('x1', 0) + .attr('y1', nodeTypeHeight + 50 + i * 20) + .attr('x2', 20) + .attr('y2', nodeTypeHeight + 50 + i * 20) + .attr('stroke', rel.color) + .attr('stroke-width', rel.type === 'contains' ? 1 : 2); + + legendGroup + .append('text') + .attr('x', 30) + .attr('y', nodeTypeHeight + 54 + i * 20) + .text(rel.label) + .style('font-size', '12px') + .style('fill', '#333'); + }); }; return ( diff --git a/pyproject.toml b/pyproject.toml index 8a84f96..2203d36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ ] dependencies = [ "pathspec>=0.11.0", + "astroid>=2.15.0", ] [project.optional-dependencies] diff --git a/src/repo_visualizer/analyzer.py b/src/repo_visualizer/analyzer.py index 1055ed6..b0c7659 100644 --- a/src/repo_visualizer/analyzer.py +++ b/src/repo_visualizer/analyzer.py @@ -13,6 +13,7 @@ from datetime import datetime from typing import Any, Dict, List, Optional, Set, Tuple +import astroid import pathspec from .schema import ( @@ -890,57 +891,11 @@ def _extract_file_relationships( extension: File extension """ if extension == "py": - # Extract Python imports - more comprehensive patterns - import_patterns = [ - # Standard imports - r"import\s+([\w\.]+(?:\s*,\s*[\w\.]+)*)", - # From imports with specific items - r"from\s+([\w\.]+)\s+import\s+(?:[\w\*]+(?:\s*,\s*[\w\*]+)*|\((?:[\w\*]+(?:\s*,\s*[\w\*]+)*)\))", - # Relative imports - r"from\s+(\.*[\w\.]*)\s+import", - # Import with aliases - r"import\s+([\w\.]+)\s+as\s+\w+", - # From import with alias - r"from\s+([\w\.]+)\s+import\s+[\w\*]+\s+as\s+\w+", - ] - - # Process each import pattern - for pattern in import_patterns: - try: - matches = re.finditer(pattern, content) - for match in matches: - raw_modules = match.group(1) - - # Handle comma-separated imports - if "," in raw_modules and "from" not in match.group(0): - modules = [m.strip() for m in raw_modules.split(",")] - else: - modules = [raw_modules] - - for module in modules: - # Try to resolve the import to a file in the repository - import_paths = self._resolve_python_import( - module, file_path - ) - for import_path in import_paths: - if import_path and import_path in self.file_ids: - relationship: Relationship = { - "source": file_path, - "target": import_path, - "type": "import", - } - # Check if this relationship already exists - if not any( - r["source"] == relationship["source"] - and r["target"] == relationship["target"] - and r["type"] == relationship["type"] - for r in self.relationships - ): - self.relationships.append(relationship) - except Exception as e: - print( - f"Error extracting Python relationships from {file_path}: {e}" - ) + try: + # Use astroid for Python code analysis + self._analyze_python_imports_with_astroid(content, file_path) + except Exception as e: + print(f"Error extracting Python relationships from {file_path}: {e}") # Look for function calls between modules self._extract_python_function_calls(content, file_path) @@ -983,6 +938,235 @@ def _extract_file_relationships( except Exception as e: print(f"Error extracting JS/TS relationships from {file_path}: {e}") + def _analyze_python_imports_with_astroid( + self, content: str, file_path: str + ) -> None: + """ + Analyze Python imports using astroid for robust AST-based parsing. + + Args: + content: Python file content + file_path: Path to the file being analyzed + """ + try: + # Parse the Python code into an AST + module = astroid.parse(content, file_path) + + # Walk through the AST to find import nodes + for node in module.body: + # Handle 'import x' statements + if isinstance(node, astroid.Import): + for name, _alias in node.names: + # Resolve the module to a file path + import_paths = self._resolve_python_import(name, file_path) + self._create_file_import_relationships(file_path, import_paths) + + # Handle 'from x import y' statements + elif isinstance(node, astroid.ImportFrom): + module_name = node.modname or "" + level = node.level or 0 # For relative imports + + # Handle relative imports + if level > 0: + # For relative imports, prepend dots to the module name + relative_prefix = "." * level + if module_name: + full_module_name = f"{relative_prefix}{module_name}" + else: + full_module_name = relative_prefix + else: + full_module_name = module_name + + # Resolve the module to file paths + import_paths = self._resolve_python_import( + full_module_name, file_path + ) + + # Create file-to-file import relationships + self._create_file_import_relationships(file_path, import_paths) + + # Process specific component imports + for import_path in import_paths: + if import_path and import_path in self.file_ids: + # Get components from the target file + target_components = self._get_file_components(import_path) + + # Process each imported name + for name, _alias in node.names: + if name == "*": # Star import + # Import all components from the module + self._create_component_import_relationships( + file_path, + import_path, + target_components, + is_star_import=True, + ) + else: + # Import specific components + self._create_component_import_relationships( + file_path, + import_path, + target_components, + component_names=[name], + ) + except Exception as e: + print(f"Error in astroid analysis for {file_path}: {e}") + + def _create_file_import_relationships( + self, source_file: str, import_paths: List[str] + ) -> None: + """ + Create file-to-file import relationships. + + Args: + source_file: The file doing the importing + import_paths: List of resolved file paths being imported + """ + for import_path in import_paths: + if import_path and import_path in self.file_ids: + relationship: Relationship = { + "source": source_file, + "target": import_path, + "type": "import", + } + # Check if this relationship already exists + if not any( + r["source"] == relationship["source"] + and r["target"] == relationship["target"] + and r["type"] == relationship["type"] + for r in self.relationships + ): + self.relationships.append(relationship) + + def _get_file_components(self, file_path: str) -> List[Component]: + """ + Get all components defined in a file. + + Args: + file_path: Path to the file + + Returns: + List of components in the file + """ + for file in self.data.get("files", []): + if file["id"] == file_path: + return file.get("components", []) + return [] + + def _create_component_import_relationships( + self, + source_file: str, + target_file: str, + target_components: List[Component], + component_names: Optional[List[str]] = None, + is_star_import: bool = False, + ) -> None: + """ + Create relationships between a file and imported components. + + Args: + source_file: The file doing the importing + target_file: The file being imported from + target_components: List of components in the target file + component_names: List of specific component names being imported + is_star_import: Whether this is a star import (import *) + """ + if not target_components: + return + + if is_star_import: + # For star imports, create relationships with all components + for component in target_components: + self._add_component_import_relationship(source_file, component["id"]) + elif component_names: + # For specific imports, create relationships only with named components + for component in target_components: + if component["name"] in component_names: + self._add_component_import_relationship( + source_file, component["id"] + ) + + def _add_component_import_relationship( + self, source_file: str, component_id: str + ) -> None: + """ + Add a component import relationship if it doesn't already exist. + + Args: + source_file: The importing file + component_id: ID of the component being imported + """ + relationship: Relationship = { + "source": source_file, + "target": component_id, + "type": "imports_component", + } + + # Check if this relationship already exists + if not any( + r["source"] == relationship["source"] + and r["target"] == relationship["target"] + and r["type"] == relationship["type"] + for r in self.relationships + ): + self.relationships.append(relationship) + + def _extract_component_imports( + self, import_statement: str, source_file: str, target_file: str + ) -> None: + """ + Compatibility method for tests - emulates the old regex-based import extraction. + + Args: + import_statement: The full import statement string + source_file: The file doing the importing + target_file: The file being imported from + """ + # Get components in the target file + target_components = self._get_file_components(target_file) + + if not target_components: + return + + # Handle star imports + if "*" in import_statement and "import *" in import_statement: + # Import all components + self._create_component_import_relationships( + source_file, target_file, target_components, is_star_import=True + ) + return + + # Extract component names from the import statement + component_names = [] + + # Handle different import patterns + if "from" in import_statement and "import" in import_statement: + # Extract the part after "import" + after_import = import_statement.split("import")[1].strip() + + # Handle parenthesized imports: from utils import (Class, Function) + if after_import.startswith("(") and ")" in after_import: + after_import = after_import.strip("()") + + # Split by comma to get individual components + for item in after_import.split(","): + item = item.strip() + if " as " in item: + # For aliased imports, get the original name + component_name = item.split(" as ")[0].strip() + component_names.append(component_name) + elif item: # Skip empty items + component_names.append(item) + + # Create relationships for the extracted component names + if component_names: + self._create_component_import_relationships( + source_file, + target_file, + target_components, + component_names=component_names, + ) + def _extract_python_function_calls(self, content: str, file_path: str) -> None: """ Extract Python function calls between components. diff --git a/tests/test_relationships.py b/tests/test_relationships.py index 74a9b16..d2b29cf 100644 --- a/tests/test_relationships.py +++ b/tests/test_relationships.py @@ -3,6 +3,7 @@ from unittest.mock import patch from repo_visualizer.analyzer import RepositoryAnalyzer +from repo_visualizer.schema import create_empty_schema class TestRepositoryRelationships: @@ -333,3 +334,102 @@ def test_duplicate_relationship_removal(self, mock_isdir): assert len(import_rels) == 2 # Two unique imports assert len(contains_rels) == 1 # One unique contains + + @patch("os.path.isdir") + def test_component_import_extraction(self, mock_isdir): + """Test extraction of specific component imports.""" + mock_isdir.return_value = True + + # Create analyzer with mock files and components + analyzer = RepositoryAnalyzer("/fake/repo") + + # Create target file with components that can be imported + target_file = { + "id": "utils.py", + "path": "utils.py", + "name": "utils.py", + "type": "file", + "components": [ + { + "id": "utils.py:HelperClass", + "name": "HelperClass", + "type": "class", + "components": [], + }, + { + "id": "utils.py:utility_function", + "name": "utility_function", + "type": "function", + "components": [], + }, + { + "id": "utils.py:CONSTANT", + "name": "CONSTANT", + "type": "variable", + "components": [], + }, + ], + } + + # Initialize with the correct type structure using the proper schema function + analyzer.data = create_empty_schema() + analyzer.data["files"] = [target_file] + analyzer.file_ids = {"utils.py"} + + # Test different import styles + import_statements = [ + # Regular component import + "from utils import HelperClass, utility_function", + # Import with alias + "from utils import HelperClass as HC, CONSTANT as C", + # Parenthesized import + "from utils import (HelperClass, utility_function)", + # Star import + "from utils import *", + ] + + # Test each import style + for import_statement in import_statements: + # Clear existing relationships + analyzer.relationships = [] + + # Call the method directly + analyzer._extract_component_imports(import_statement, "main.py", "utils.py") + + # Check for component import relationships + if "*" in import_statement: + # Star import should import all components + assert len(analyzer.relationships) == 3 + else: + # Should have relationships for the specific components + component_imports = [ + r + for r in analyzer.relationships + if r["type"] == "imports_component" + ] + + # Check that HelperClass was imported + assert any( + r["source"] == "main.py" + and r["target"] == "utils.py:HelperClass" + and r["type"] == "imports_component" + for r in component_imports + ) + + # utility_function should be imported in regular & parenthesized imports + if "utility_function" in import_statement: + assert any( + r["source"] == "main.py" + and r["target"] == "utils.py:utility_function" + and r["type"] == "imports_component" + for r in component_imports + ) + + # CONSTANT should be imported in the aliased import + if "CONSTANT" in import_statement: + assert any( + r["source"] == "main.py" + and r["target"] == "utils.py:CONSTANT" + and r["type"] == "imports_component" + for r in component_imports + )