hanzo-mcp 0.7.7__py3-none-any.whl → 0.8.1__py3-none-any.whl

hanzo_mcp/tools/vector/ast_analyzer.py

@@ -1,15 +1,15 @@
 """AST analysis and symbol extraction for code understanding."""
 
 import ast
-import json
-from pathlib import Path
-from typing import Dict, List, Any, Optional, Set, Tuple
-from dataclasses import dataclass, asdict
 import hashlib
+from typing import Any, Dict, List, Optional
+from pathlib import Path
+from dataclasses import asdict, dataclass
 
 try:
-    import tree_sitter_python as tspython
     import tree_sitter
+    import tree_sitter_python as tspython
+
     TREE_SITTER_AVAILABLE = True
 except ImportError:
     TREE_SITTER_AVAILABLE = False
@@ -18,6 +18,7 @@ except ImportError:
 @dataclass
 class Symbol:
     """Represents a code symbol (function, class, variable, etc.)."""
+
     name: str
     type: str  # function, class, variable, import, etc.
     file_path: str
@@ -30,24 +31,25 @@ class Symbol:
     docstring: Optional[str] = None
     signature: Optional[str] = None
     references: List[str] = None  # Files that reference this symbol
-    
+
     def __post_init__(self):
         if self.references is None:
             self.references = []
 
 
-@dataclass 
+@dataclass
 class ASTNode:
     """Represents an AST node with metadata."""
+
     type: str
     name: Optional[str]
     line_start: int
     line_end: int
     column_start: int
     column_end: int
-    children: List['ASTNode'] = None
+    children: List["ASTNode"] = None
     parent: Optional[str] = None
-    
+
     def __post_init__(self):
         if self.children is None:
             self.children = []
@@ -56,6 +58,7 @@ class ASTNode:
 @dataclass
 class FileAST:
     """Complete AST representation of a file."""
+
     file_path: str
     file_hash: str
     language: str
@@ -64,7 +67,7 @@ class FileAST:
     imports: List[str]
     exports: List[str]
     dependencies: List[str]  # Files this file depends on
-    
+
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary for storage."""
         return {
@@ -81,182 +84,198 @@ class FileAST:
 
 class ASTAnalyzer:
     """Analyzes code files and extracts AST information and symbols."""
-    
+
     def __init__(self):
         """Initialize the AST analyzer."""
         self.parsers = {}
         self._setup_parsers()
-    
+
     def _setup_parsers(self):
         """Set up tree-sitter parsers for different languages."""
         if TREE_SITTER_AVAILABLE:
             try:
                 # Python parser
-                self.parsers['python'] = tree_sitter.Language(tspython.language())
+                self.parsers["python"] = tree_sitter.Language(tspython.language())
             except Exception as e:
                 import logging
+
                 logger = logging.getLogger(__name__)
                 logger.warning(f"Could not initialize Python parser: {e}")
-    
+
     def analyze_file(self, file_path: str) -> Optional[FileAST]:
         """Analyze a file and extract AST information and symbols.
-        
+
         Args:
             file_path: Path to the file to analyze
-            
+
         Returns:
             FileAST object with extracted information, or None if analysis fails
         """
         path = Path(file_path)
         if not path.exists():
             return None
-        
+
         # Determine language
         language = self._detect_language(path)
         if not language:
             return None
-        
+
         try:
             # Read file content
-            content = path.read_text(encoding='utf-8')
+            content = path.read_text(encoding="utf-8")
             file_hash = hashlib.sha256(content.encode()).hexdigest()
-            
+
             # Extract symbols and AST
-            if language == 'python':
+            if language == "python":
                 return self._analyze_python_file(file_path, content, file_hash)
             else:
                 # Generic analysis for other languages
-                return self._analyze_generic_file(file_path, content, file_hash, language)
-                
+                return self._analyze_generic_file(
+                    file_path, content, file_hash, language
+                )
+
         except Exception as e:
             import logging
+
             logger = logging.getLogger(__name__)
             logger.error(f"Error analyzing file {file_path}: {e}")
             return None
-    
+
     def _detect_language(self, path: Path) -> Optional[str]:
         """Detect programming language from file extension."""
         extension = path.suffix.lower()
-        
+
         language_map = {
-            '.py': 'python',
-            '.js': 'javascript',
-            '.ts': 'typescript',
-            '.jsx': 'javascript',
-            '.tsx': 'typescript',
-            '.java': 'java',
-            '.cpp': 'cpp',
-            '.c': 'c',
-            '.h': 'c',
-            '.hpp': 'cpp',
-            '.rs': 'rust',
-            '.go': 'go',
-            '.rb': 'ruby',
-            '.php': 'php',
-            '.cs': 'csharp',
-            '.swift': 'swift',
-            '.kt': 'kotlin',
-            '.scala': 'scala',
-            '.clj': 'clojure',
-            '.hs': 'haskell',
-            '.ml': 'ocaml',
-            '.elm': 'elm',
-            '.dart': 'dart',
-            '.lua': 'lua',
-            '.r': 'r',
-            '.m': 'objective-c',
-            '.mm': 'objective-cpp',
+            ".py": "python",
+            ".js": "javascript",
+            ".ts": "typescript",
+            ".jsx": "javascript",
+            ".tsx": "typescript",
+            ".java": "java",
+            ".cpp": "cpp",
+            ".c": "c",
+            ".h": "c",
+            ".hpp": "cpp",
+            ".rs": "rust",
+            ".go": "go",
+            ".rb": "ruby",
+            ".php": "php",
+            ".cs": "csharp",
+            ".swift": "swift",
+            ".kt": "kotlin",
+            ".scala": "scala",
+            ".clj": "clojure",
+            ".hs": "haskell",
+            ".ml": "ocaml",
+            ".elm": "elm",
+            ".dart": "dart",
+            ".lua": "lua",
+            ".r": "r",
+            ".m": "objective-c",
+            ".mm": "objective-cpp",
         }
-        
+
         return language_map.get(extension)
-    
-    def _analyze_python_file(self, file_path: str, content: str, file_hash: str) -> FileAST:
+
+    def _analyze_python_file(
+        self, file_path: str, content: str, file_hash: str
+    ) -> FileAST:
         """Analyze Python file using both AST and tree-sitter."""
         symbols = []
         ast_nodes = []
         imports = []
         exports = []
         dependencies = []
-        
+
         try:
             # Parse with Python AST
             tree = ast.parse(content)
-            
+
             # Extract symbols using AST visitor
             visitor = PythonSymbolExtractor(file_path)
             visitor.visit(tree)
-            
+
             symbols.extend(visitor.symbols)
             imports.extend(visitor.imports)
             exports.extend(visitor.exports)
             dependencies.extend(visitor.dependencies)
-            
+
             # If tree-sitter is available, get more detailed AST
-            if TREE_SITTER_AVAILABLE and 'python' in self.parsers:
-                parser = tree_sitter.Parser(self.parsers['python'])
+            if TREE_SITTER_AVAILABLE and "python" in self.parsers:
+                parser = tree_sitter.Parser(self.parsers["python"])
                 ts_tree = parser.parse(content.encode())
                 ast_nodes = self._extract_tree_sitter_nodes(ts_tree.root_node, content)
-            
+
         except SyntaxError as e:
             import logging
+
             logger = logging.getLogger(__name__)
             logger.error(f"Syntax error in {file_path}: {e}")
         except Exception as e:
             import logging
+
             logger = logging.getLogger(__name__)
             logger.error(f"Error parsing Python file {file_path}: {e}")
-        
+
         return FileAST(
             file_path=file_path,
             file_hash=file_hash,
-            language='python',
+            language="python",
             symbols=symbols,
             ast_nodes=ast_nodes,
             imports=imports,
             exports=exports,
             dependencies=dependencies,
         )
-    
-    def _analyze_generic_file(self, file_path: str, content: str, file_hash: str, language: str) -> FileAST:
+
+    def _analyze_generic_file(
+        self, file_path: str, content: str, file_hash: str, language: str
+    ) -> FileAST:
         """Generic analysis for non-Python files."""
         # For now, just basic line-based analysis
         # Could be enhanced with language-specific parsers
-        
+
         symbols = []
         ast_nodes = []
         imports = []
         exports = []
         dependencies = []
-        
+
         # Basic pattern matching for common constructs
-        lines = content.split('\n')
+        lines = content.split("\n")
         for i, line in enumerate(lines, 1):
             line = line.strip()
-            
+
             # Basic function detection (works for many C-style languages)
-            if language in ['javascript', 'typescript', 'java', 'cpp', 'c']:
-                if 'function ' in line or line.startswith('def ') or ' function(' in line:
+            if language in ["javascript", "typescript", "java", "cpp", "c"]:
+                if (
+                    "function " in line
+                    or line.startswith("def ")
+                    or " function(" in line
+                ):
                     # Extract function name
                     parts = line.split()
                     for j, part in enumerate(parts):
-                        if part == 'function' and j + 1 < len(parts):
-                            func_name = parts[j + 1].split('(')[0]
-                            symbols.append(Symbol(
-                                name=func_name,
-                                type='function',
-                                file_path=file_path,
-                                line_start=i,
-                                line_end=i,
-                                column_start=0,
-                                column_end=len(line),
-                                scope='global',
-                            ))
+                        if part == "function" and j + 1 < len(parts):
+                            func_name = parts[j + 1].split("(")[0]
+                            symbols.append(
+                                Symbol(
+                                    name=func_name,
+                                    type="function",
+                                    file_path=file_path,
+                                    line_start=i,
+                                    line_end=i,
+                                    column_start=0,
+                                    column_end=len(line),
+                                    scope="global",
+                                )
+                            )
                             break
-            
+
             # Basic import detection
-            if 'import ' in line or '#include ' in line or 'require(' in line:
+            if "import " in line or "#include " in line or "require(" in line:
                 imports.append(line)
-        
+
         return FileAST(
             file_path=file_path,
             file_hash=file_hash,
@@ -267,23 +286,27 @@ class ASTAnalyzer:
             exports=exports,
             dependencies=dependencies,
         )
-    
+
     def _extract_tree_sitter_nodes(self, node, content: str) -> List[ASTNode]:
         """Extract AST nodes from tree-sitter parse tree."""
         nodes = []
-        
+
         def traverse(ts_node, parent_name=None):
             node_name = None
-            
+
             # Try to extract node name for named nodes
-            if ts_node.type in ['function_definition', 'class_definition', 'identifier']:
+            if ts_node.type in [
+                "function_definition",
+                "class_definition",
+                "identifier",
+            ]:
                 for child in ts_node.children:
-                    if child.type == 'identifier':
+                    if child.type == "identifier":
                         start_byte = child.start_byte
                         end_byte = child.end_byte
                         node_name = content[start_byte:end_byte]
                         break
-            
+
             ast_node = ASTNode(
                 type=ts_node.type,
                 name=node_name,
@@ -293,148 +316,158 @@ class ASTAnalyzer:
                 column_end=ts_node.end_point[1],
                 parent=parent_name,
             )
-            
+
             nodes.append(ast_node)
-            
+
             # Recursively process children
             for child in ts_node.children:
                 traverse(child, node_name or parent_name)
-        
+
         traverse(node)
         return nodes
 
 
 class PythonSymbolExtractor(ast.NodeVisitor):
     """AST visitor for extracting Python symbols."""
-    
+
     def __init__(self, file_path: str):
         self.file_path = file_path
         self.symbols = []
         self.imports = []
         self.exports = []
         self.dependencies = []
-        self.scope_stack = ['global']
-    
+        self.scope_stack = ["global"]
+
     def visit_FunctionDef(self, node):
         """Visit function definitions."""
-        scope = '.'.join(self.scope_stack)
+        scope = ".".join(self.scope_stack)
         parent = self.scope_stack[-1] if len(self.scope_stack) > 1 else None
-        
+
         # Extract docstring
         docstring = None
-        if (node.body and isinstance(node.body[0], ast.Expr) and 
-            isinstance(node.body[0].value, ast.Constant) and 
-            isinstance(node.body[0].value.value, str)):
+        if (
+            node.body
+            and isinstance(node.body[0], ast.Expr)
+            and isinstance(node.body[0].value, ast.Constant)
+            and isinstance(node.body[0].value.value, str)
+        ):
             docstring = node.body[0].value.value
-        
+
         # Create function signature
         args = [arg.arg for arg in node.args.args]
         signature = f"{node.name}({', '.join(args)})"
-        
+
         symbol = Symbol(
             name=node.name,
-            type='function',
+            type="function",
             file_path=self.file_path,
             line_start=node.lineno,
             line_end=node.end_lineno or node.lineno,
             column_start=node.col_offset,
             column_end=node.end_col_offset or 0,
             scope=scope,
-            parent=parent if parent != 'global' else None,
+            parent=parent if parent != "global" else None,
             docstring=docstring,
             signature=signature,
         )
-        
+
         self.symbols.append(symbol)
-        
+
         # Enter function scope
         self.scope_stack.append(node.name)
         self.generic_visit(node)
         self.scope_stack.pop()
-    
+
     def visit_AsyncFunctionDef(self, node):
         """Visit async function definitions."""
         self.visit_FunctionDef(node)  # Same logic
-    
+
     def visit_ClassDef(self, node):
         """Visit class definitions."""
-        scope = '.'.join(self.scope_stack)
+        scope = ".".join(self.scope_stack)
         parent = self.scope_stack[-1] if len(self.scope_stack) > 1 else None
-        
+
         # Extract docstring
         docstring = None
-        if (node.body and isinstance(node.body[0], ast.Expr) and 
-            isinstance(node.body[0].value, ast.Constant) and 
-            isinstance(node.body[0].value.value, str)):
+        if (
+            node.body
+            and isinstance(node.body[0], ast.Expr)
+            and isinstance(node.body[0].value, ast.Constant)
+            and isinstance(node.body[0].value.value, str)
+        ):
             docstring = node.body[0].value.value
-        
+
         # Extract base classes
         bases = [self._get_name(base) for base in node.bases]
-        signature = f"class {node.name}({', '.join(bases)})" if bases else f"class {node.name}"
-        
+        signature = (
+            f"class {node.name}({', '.join(bases)})" if bases else f"class {node.name}"
+        )
+
         symbol = Symbol(
             name=node.name,
-            type='class',
+            type="class",
             file_path=self.file_path,
             line_start=node.lineno,
             line_end=node.end_lineno or node.lineno,
             column_start=node.col_offset,
             column_end=node.end_col_offset or 0,
             scope=scope,
-            parent=parent if parent != 'global' else None,
+            parent=parent if parent != "global" else None,
             docstring=docstring,
             signature=signature,
         )
-        
+
         self.symbols.append(symbol)
-        
+
         # Enter class scope
         self.scope_stack.append(node.name)
         self.generic_visit(node)
         self.scope_stack.pop()
-    
+
     def visit_Import(self, node):
         """Visit import statements."""
         for alias in node.names:
             import_name = alias.name
             self.imports.append(import_name)
-            if '.' not in import_name:  # Top-level module
+            if "." not in import_name:  # Top-level module
                 self.dependencies.append(import_name)
-    
+
     def visit_ImportFrom(self, node):
         """Visit from...import statements."""
         if node.module:
             self.imports.append(node.module)
-            if '.' not in node.module:  # Top-level module
+            if "." not in node.module:  # Top-level module
                 self.dependencies.append(node.module)
-        
+
         for alias in node.names:
-            if alias.name != '*':
-                import_item = f"{node.module}.{alias.name}" if node.module else alias.name
+            if alias.name != "*":
+                import_item = (
+                    f"{node.module}.{alias.name}" if node.module else alias.name
+                )
                 self.imports.append(import_item)
-    
+
     def visit_Assign(self, node):
         """Visit variable assignments."""
-        scope = '.'.join(self.scope_stack)
+        scope = ".".join(self.scope_stack)
         parent = self.scope_stack[-1] if len(self.scope_stack) > 1 else None
-        
+
         for target in node.targets:
             if isinstance(target, ast.Name):
                 symbol = Symbol(
                     name=target.id,
-                    type='variable',
+                    type="variable",
                     file_path=self.file_path,
                     line_start=node.lineno,
                     line_end=node.end_lineno or node.lineno,
                     column_start=node.col_offset,
                     column_end=node.end_col_offset or 0,
                     scope=scope,
-                    parent=parent if parent != 'global' else None,
+                    parent=parent if parent != "global" else None,
                 )
                 self.symbols.append(symbol)
-        
+
         self.generic_visit(node)
-    
+
     def _get_name(self, node):
         """Extract name from AST node."""
         if isinstance(node, ast.Name):
@@ -454,14 +487,14 @@ def create_symbol_embedding_text(symbol: Symbol) -> str:
         f"Type: {symbol.type}",
         f"Scope: {symbol.scope}",
     ]
-    
+
     if symbol.parent:
         parts.append(f"Parent: {symbol.parent}")
-    
+
     if symbol.signature:
         parts.append(f"Signature: {symbol.signature}")
-    
+
     if symbol.docstring:
         parts.append(f"Documentation: {symbol.docstring}")
-    
-    return " | ".join(parts)
+
+    return " | ".join(parts)