hanzo-mcp 0.6.12__py3-none-any.whl → 0.7.0__py3-none-any.whl

hanzo_mcp/tools/filesystem/ast_multi_edit.py

@@ -0,0 +1,562 @@
+"""AST-aware multi-edit tool using treesitter for accurate code modifications."""
+
+import os
+import json
+from typing import List, Dict, Any, Optional, Tuple, Set
+from pathlib import Path
+from dataclasses import dataclass
+from collections import defaultdict
+
+from hanzo_mcp.tools.common.base import BaseTool
+from hanzo_mcp.tools.common.decorators import with_context_normalization
+from hanzo_mcp.tools.common.paginated_response import AutoPaginatedResponse
+from hanzo_mcp.types import MCPResourceDocument
+
+try:
+    import tree_sitter
+    import tree_sitter_python
+    import tree_sitter_javascript
+    import tree_sitter_typescript
+    import tree_sitter_go
+    import tree_sitter_rust
+    import tree_sitter_java
+    import tree_sitter_cpp
+    TREESITTER_AVAILABLE = True
+except ImportError:
+    TREESITTER_AVAILABLE = False
+
+
+@dataclass
+class ASTMatch:
+    """Represents an AST match with context."""
+    file_path: str
+    line_start: int
+    line_end: int
+    column_start: int
+    column_end: int
+    node_type: str
+    text: str
+    parent_context: Optional[str] = None
+    semantic_context: Optional[str] = None
+
+
+@dataclass 
+class EditOperation:
+    """Enhanced edit operation with AST awareness."""
+    old_string: str
+    new_string: str
+    node_types: Optional[List[str]] = None  # Restrict to specific AST node types
+    semantic_match: bool = False  # Use semantic matching
+    expect_count: Optional[int] = None  # Expected number of matches
+    context_lines: int = 5  # Lines of context for uniqueness
+
+
+class ASTMultiEdit(BaseTool):
+    """Multi-edit tool with AST awareness and automatic reference finding."""
+    
+    name = "ast_multi_edit"
+    description = """Enhanced multi-edit with AST awareness and reference finding.
+    
+    Features:
+    - AST-based search for accurate matches
+    - Automatic reference finding across codebase
+    - Semantic matching (find all usages of a symbol)
+    - Result pagination to avoid token limits
+    - Context-aware replacements
+    
+    Examples:
+    1. Rename a function and all its calls:
+       ast_multi_edit("file.py", [
+           {"old_string": "oldFunc", "new_string": "newFunc", "semantic_match": true}
+       ])
+    
+    2. Update specific node types only:
+       ast_multi_edit("file.go", [
+           {"old_string": "StopTracking", "new_string": "StopTrackingWithContext",
+            "node_types": ["call_expression"]}
+       ])
+    """
+    
+    def __init__(self):
+        super().__init__()
+        self.parsers = {}
+        self.languages = {}
+        
+        if TREESITTER_AVAILABLE:
+            self._init_parsers()
+    
+    def _init_parsers(self):
+        """Initialize treesitter parsers for supported languages."""
+        language_mapping = {
+            '.py': (tree_sitter_python, 'python'),
+            '.js': (tree_sitter_javascript, 'javascript'),
+            '.jsx': (tree_sitter_javascript, 'javascript'),
+            '.ts': (tree_sitter_typescript.typescript, 'typescript'),
+            '.tsx': (tree_sitter_typescript.tsx, 'tsx'),
+            '.go': (tree_sitter_go, 'go'),
+            '.rs': (tree_sitter_rust, 'rust'),
+            '.java': (tree_sitter_java, 'java'),
+            '.cpp': (tree_sitter_cpp, 'cpp'),
+            '.cc': (tree_sitter_cpp, 'cpp'),
+            '.cxx': (tree_sitter_cpp, 'cpp'),
+            '.h': (tree_sitter_cpp, 'cpp'),
+            '.hpp': (tree_sitter_cpp, 'cpp'),
+        }
+        
+        for ext, (module, name) in language_mapping.items():
+            try:
+                parser = tree_sitter.Parser()
+                if hasattr(module, 'language'):
+                    parser.set_language(module.language())
+                else:
+                    # For older tree-sitter bindings
+                    lang = tree_sitter.Language(module.language(), name)
+                    parser.set_language(lang)
+                self.parsers[ext] = parser
+                self.languages[ext] = name
+            except Exception as e:
+                print(f"Failed to initialize parser for {ext}: {e}")
+    
+    def _get_parser(self, file_path: str) -> Optional[tree_sitter.Parser]:
+        """Get parser for file type."""
+        ext = Path(file_path).suffix.lower()
+        return self.parsers.get(ext)
+    
+    def _parse_file(self, file_path: str, content: str) -> Optional[tree_sitter.Tree]:
+        """Parse file content into AST."""
+        parser = self._get_parser(file_path)
+        if not parser:
+            return None
+        
+        return parser.parse(bytes(content, 'utf-8'))
+    
+    def _find_references(self, 
+                        symbol: str, 
+                        file_path: str,
+                        project_root: Optional[str] = None) -> List[ASTMatch]:
+        """Find all references to a symbol across the project."""
+        matches = []
+        
+        if not project_root:
+            project_root = self._find_project_root(file_path)
+        
+        # Get language-specific reference patterns
+        patterns = self._get_reference_patterns(symbol, file_path)
+        
+        # Search across all relevant files
+        for pattern in patterns:
+            # Use grep_ast tool for efficient AST-aware search
+            results = self._search_with_ast(pattern, project_root)
+            matches.extend(results)
+        
+        return matches
+    
+    def _get_reference_patterns(self, symbol: str, file_path: str) -> List[Dict[str, Any]]:
+        """Get language-specific patterns for finding references."""
+        ext = Path(file_path).suffix.lower()
+        lang = self.languages.get(ext, 'generic')
+        
+        patterns = []
+        
+        if lang == 'go':
+            # Go specific patterns
+            patterns.extend([
+                # Function calls
+                {"query": f"(call_expression function: (identifier) @func (#eq? @func \"{symbol}\"))", "type": "call"},
+                # Method calls  
+                {"query": f"(call_expression function: (selector_expression field: (field_identifier) @method (#eq? @method \"{symbol}\")))", "type": "method_call"},
+                # Function declarations
+                {"query": f"(function_declaration name: (identifier) @name (#eq? @name \"{symbol}\"))", "type": "declaration"},
+                # Type references
+                {"query": f"(type_identifier) @type (#eq? @type \"{symbol}\")", "type": "type_ref"},
+            ])
+        elif lang in ['javascript', 'typescript', 'tsx']:
+            patterns.extend([
+                # Function calls
+                {"query": f"(call_expression function: (identifier) @func (#eq? @func \"{symbol}\"))", "type": "call"},
+                # Method calls
+                {"query": f"(call_expression function: (member_expression property: (property_identifier) @prop (#eq? @prop \"{symbol}\")))", "type": "method_call"},
+                # Function declarations
+                {"query": f"(function_declaration name: (identifier) @name (#eq? @name \"{symbol}\"))", "type": "declaration"},
+                # Variable declarations
+                {"query": f"(variable_declarator name: (identifier) @var (#eq? @var \"{symbol}\"))", "type": "variable"},
+            ])
+        elif lang == 'python':
+            patterns.extend([
+                # Function calls
+                {"query": f"(call function: (identifier) @func (#eq? @func \"{symbol}\"))", "type": "call"},
+                # Method calls
+                {"query": f"(call function: (attribute attribute: (identifier) @attr (#eq? @attr \"{symbol}\")))", "type": "method_call"},
+                # Function definitions
+                {"query": f"(function_definition name: (identifier) @name (#eq? @name \"{symbol}\"))", "type": "declaration"},
+                # Class definitions
+                {"query": f"(class_definition name: (identifier) @name (#eq? @name \"{symbol}\"))", "type": "class"},
+            ])
+        else:
+            # Generic patterns
+            patterns.append({"query": symbol, "type": "text"})
+        
+        return patterns
+    
+    def _search_with_ast(self, pattern: Dict[str, Any], root: str) -> List[ASTMatch]:
+        """Search using AST patterns."""
+        matches = []
+        
+        # This would integrate with grep_ast tool
+        # For now, simulate the search
+        import glob
+        
+        for file_path in glob.glob(f"{root}/**/*.*", recursive=True):
+            if self._should_skip_file(file_path):
+                continue
+                
+            try:
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+                
+                tree = self._parse_file(file_path, content)
+                if tree and pattern["type"] != "text":
+                    # Use treesitter query
+                    matches.extend(self._query_ast(tree, pattern, file_path, content))
+                else:
+                    # Fallback to text search
+                    matches.extend(self._text_search(content, pattern["query"], file_path))
+                    
+            except Exception:
+                continue
+        
+        return matches
+    
+    def _query_ast(self, 
+                   tree: tree_sitter.Tree, 
+                   pattern: Dict[str, Any], 
+                   file_path: str,
+                   content: str) -> List[ASTMatch]:
+        """Query AST with treesitter pattern."""
+        matches = []
+        
+        try:
+            # Get language for query
+            lang_name = self.languages.get(Path(file_path).suffix.lower())
+            if not lang_name:
+                return matches
+            
+            # Execute query
+            query = tree_sitter.Query(pattern["query"], lang_name)
+            captures = query.captures(tree.root_node)
+            
+            lines = content.split('\n')
+            
+            for node, name in captures:
+                match = ASTMatch(
+                    file_path=file_path,
+                    line_start=node.start_point[0] + 1,
+                    line_end=node.end_point[0] + 1,
+                    column_start=node.start_point[1],
+                    column_end=node.end_point[1],
+                    node_type=node.type,
+                    text=content[node.start_byte:node.end_byte],
+                    parent_context=self._get_parent_context(node, content),
+                    semantic_context=pattern["type"]
+                )
+                matches.append(match)
+                
+        except Exception as e:
+            # Fallback to simple search
+            pass
+        
+        return matches
+    
+    def _get_parent_context(self, node: tree_sitter.Node, content: str) -> Optional[str]:
+        """Get parent context for better understanding."""
+        parent = node.parent
+        if parent:
+            # Get parent function/class name
+            if parent.type in ['function_declaration', 'function_definition', 'method_definition']:
+                for child in parent.children:
+                    if child.type == 'identifier':
+                        return f"function: {content[child.start_byte:child.end_byte]}"
+            elif parent.type in ['class_declaration', 'class_definition']:
+                for child in parent.children:
+                    if child.type == 'identifier':
+                        return f"class: {content[child.start_byte:child.end_byte]}"
+        
+        return None
+    
+    def _text_search(self, content: str, pattern: str, file_path: str) -> List[ASTMatch]:
+        """Fallback text search."""
+        matches = []
+        lines = content.split('\n')
+        
+        for i, line in enumerate(lines):
+            if pattern in line:
+                col = line.find(pattern)
+                match = ASTMatch(
+                    file_path=file_path,
+                    line_start=i + 1,
+                    line_end=i + 1,
+                    column_start=col,
+                    column_end=col + len(pattern),
+                    node_type='text',
+                    text=pattern,
+                    semantic_context='text_match'
+                )
+                matches.append(match)
+        
+        return matches
+    
+    def _should_skip_file(self, file_path: str) -> bool:
+        """Check if file should be skipped."""
+        skip_dirs = {'.git', 'node_modules', '__pycache__', '.pytest_cache', 'venv', '.env'}
+        skip_extensions = {'.pyc', '.pyo', '.so', '.dylib', '.dll', '.exe'}
+        
+        path = Path(file_path)
+        
+        # Check directories
+        for part in path.parts:
+            if part in skip_dirs:
+                return True
+        
+        # Check extensions
+        if path.suffix in skip_extensions:
+            return True
+        
+        # Check if binary
+        try:
+            with open(file_path, 'rb') as f:
+                chunk = f.read(512)
+                if b'\0' in chunk:
+                    return True
+        except:
+            return True
+        
+        return False
+    
+    def _find_project_root(self, file_path: str) -> str:
+        """Find project root by looking for markers."""
+        markers = {'.git', 'package.json', 'go.mod', 'Cargo.toml', 'pyproject.toml', 'setup.py'}
+        
+        path = Path(file_path).resolve()
+        for parent in path.parents:
+            for marker in markers:
+                if (parent / marker).exists():
+                    return str(parent)
+        
+        return str(path.parent)
+    
+    def _group_matches_by_file(self, matches: List[ASTMatch]) -> Dict[str, List[ASTMatch]]:
+        """Group matches by file for efficient editing."""
+        grouped = defaultdict(list)
+        for match in matches:
+            grouped[match.file_path].append(match)
+        return grouped
+    
+    def _create_unique_context(self, 
+                              content: str, 
+                              match: ASTMatch, 
+                              context_lines: int) -> str:
+        """Create unique context for edit identification."""
+        lines = content.split('\n')
+        
+        start_line = max(0, match.line_start - context_lines - 1)
+        end_line = min(len(lines), match.line_end + context_lines)
+        
+        context_lines = lines[start_line:end_line]
+        return '\n'.join(context_lines)
+    
+    async def run(self,
+                  file_path: str,
+                  edits: List[Dict[str, Any]],
+                  find_references: bool = False,
+                  page_size: int = 50,
+                  preview_only: bool = False,
+                  **kwargs) -> MCPResourceDocument:
+        """Execute AST-aware multi-edit operation.
+        
+        Args:
+            file_path: Primary file to edit
+            edits: List of edit operations
+            find_references: Whether to find and edit references across codebase
+            page_size: Number of results per page
+            preview_only: Show what would be changed without applying
+        """
+        
+        if not TREESITTER_AVAILABLE:
+            return self._fallback_to_basic_edit(file_path, edits)
+        
+        results = {
+            "primary_file": file_path,
+            "edits_requested": len(edits),
+            "files_analyzed": 0,
+            "matches_found": 0,
+            "edits_applied": 0,
+            "errors": [],
+            "changes": []
+        }
+        
+        # Convert edits to EditOperation objects
+        edit_ops = []
+        for edit in edits:
+            edit_ops.append(EditOperation(
+                old_string=edit["old_string"],
+                new_string=edit["new_string"],
+                node_types=edit.get("node_types"),
+                semantic_match=edit.get("semantic_match", False),
+                expect_count=edit.get("expect_count"),
+                context_lines=edit.get("context_lines", 5)
+            ))
+        
+        # Find all matches
+        all_matches = []
+        
+        # First, analyze primary file
+        try:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+            
+            tree = self._parse_file(file_path, content)
+            
+            for edit_op in edit_ops:
+                if edit_op.semantic_match and find_references:
+                    # Find all references across codebase
+                    matches = self._find_references(edit_op.old_string, file_path)
+                else:
+                    # Just search in current file
+                    if tree:
+                        pattern = {"query": edit_op.old_string, "type": "text"}
+                        matches = self._query_ast(tree, pattern, file_path, content)
+                    else:
+                        matches = self._text_search(content, edit_op.old_string, file_path)
+                
+                # Filter by node types if specified
+                if edit_op.node_types:
+                    matches = [m for m in matches if m.node_type in edit_op.node_types]
+                
+                # Check expected count
+                if edit_op.expect_count is not None and len(matches) != edit_op.expect_count:
+                    results["errors"].append({
+                        "edit": edit_op.old_string,
+                        "expected": edit_op.expect_count,
+                        "found": len(matches),
+                        "locations": [f"{m.file_path}:{m.line_start}" for m in matches[:5]]
+                    })
+                    continue
+                
+                all_matches.extend([(edit_op, match) for match in matches])
+        
+        except Exception as e:
+            results["errors"].append({
+                "file": file_path,
+                "error": str(e)
+            })
+            return MCPResourceDocument(data=results)
+        
+        results["matches_found"] = len(all_matches)
+        results["files_analyzed"] = len(set(m[1].file_path for m in all_matches))
+        
+        if preview_only:
+            # Return preview of changes
+            preview = self._generate_preview(all_matches, page_size)
+            results["preview"] = preview
+            return MCPResourceDocument(data=results)
+        
+        # Apply edits
+        changes_by_file = self._group_changes(all_matches)
+        
+        for file_path, changes in changes_by_file.items():
+            try:
+                success = await self._apply_file_changes(file_path, changes)
+                if success:
+                    results["edits_applied"] += len(changes)
+                    results["changes"].append({
+                        "file": file_path,
+                        "edits": len(changes)
+                    })
+            except Exception as e:
+                results["errors"].append({
+                    "file": file_path,
+                    "error": str(e)
+                })
+        
+        return MCPResourceDocument(data=results)
+    
+    def _group_changes(self, matches: List[Tuple[EditOperation, ASTMatch]]) -> Dict[str, List[Tuple[EditOperation, ASTMatch]]]:
+        """Group changes by file."""
+        grouped = defaultdict(list)
+        for edit_op, match in matches:
+            grouped[match.file_path].append((edit_op, match))
+        return grouped
+    
+    async def _apply_file_changes(self, 
+                                 file_path: str, 
+                                 changes: List[Tuple[EditOperation, ASTMatch]]) -> bool:
+        """Apply changes to a single file."""
+        with open(file_path, 'r', encoding='utf-8') as f:
+            content = f.read()
+        
+        # Sort changes by position (reverse order to maintain positions)
+        changes.sort(key=lambda x: (x[1].line_start, x[1].column_start), reverse=True)
+        
+        lines = content.split('\n')
+        
+        for edit_op, match in changes:
+            # Create unique context for this match
+            context = self._create_unique_context(content, match, edit_op.context_lines)
+            
+            # Apply the edit
+            if match.line_start == match.line_end:
+                # Single line edit
+                line = lines[match.line_start - 1]
+                before = line[:match.column_start]
+                after = line[match.column_end:]
+                lines[match.line_start - 1] = before + edit_op.new_string + after
+            else:
+                # Multi-line edit
+                # Remove old lines
+                del lines[match.line_start - 1:match.line_end]
+                # Insert new content
+                lines.insert(match.line_start - 1, edit_op.new_string)
+        
+        # Write back
+        with open(file_path, 'w', encoding='utf-8') as f:
+            f.write('\n'.join(lines))
+        
+        return True
+    
+    def _generate_preview(self, 
+                         matches: List[Tuple[EditOperation, ASTMatch]], 
+                         page_size: int) -> List[Dict[str, Any]]:
+        """Generate preview of changes."""
+        preview = []
+        
+        for i, (edit_op, match) in enumerate(matches[:page_size]):
+            preview.append({
+                "file": match.file_path,
+                "line": match.line_start,
+                "column": match.column_start,
+                "node_type": match.node_type,
+                "context": match.parent_context,
+                "old": edit_op.old_string,
+                "new": edit_op.new_string,
+                "semantic_type": match.semantic_context
+            })
+        
+        if len(matches) > page_size:
+            preview.append({
+                "note": f"... and {len(matches) - page_size} more matches"
+            })
+        
+        return preview
+    
+    def _fallback_to_basic_edit(self, file_path: str, edits: List[Dict[str, Any]]) -> MCPResourceDocument:
+        """Fallback to basic multi-edit when treesitter not available."""
+        # Delegate to existing multi_edit tool
+        from hanzo_mcp.tools.filesystem.multi_edit import MultiEdit
+        basic_tool = MultiEdit()
+        return basic_tool.run(file_path, edits)
+
+
+# Tool registration
+def create_ast_multi_edit_tool():
+    """Factory function to create AST multi-edit tool."""
+    return ASTMultiEdit()

hanzo_mcp/tools/filesystem/batch_search.py

@@ -45,7 +45,7 @@ class SearchType(Enum):
 
 @dataclass 
 class SearchResult:
-    """Unified search result combining different search types."""
+    """Search result combining different search types."""
     file_path: str
     line_number: Optional[int]
     content: str
@@ -112,11 +112,11 @@ class BatchSearchParams(TypedDict):
 
 @final
 class BatchSearchTool(FilesystemBaseTool):
-    """Unified search tool combining multiple search strategies."""
+    """Search tool combining multiple search strategies."""
     
     def __init__(self, permission_manager: PermissionManager, 
                  project_manager: Optional[ProjectVectorManager] = None):
-        """Initialize the unified search tool."""
+        """Initialize the search tool."""
         super().__init__(permission_manager)
         self.project_manager = project_manager
         

hanzo_mcp/tools/filesystem/diff.py

@@ -98,7 +98,7 @@ diff a.json b.json --ignore-whitespace"""
         
         # Generate diff
         if unified:
-            # Unified diff format
+            #  diff format
             diff_lines = list(difflib.unified_diff(
                 lines1,
                 lines2,
@@ -183,7 +183,7 @@ diff a.json b.json --ignore-whitespace"""
             unified: bool = True,
             context: int = 3,
             ignore_whitespace: bool = False,
-            show_line_numbers: bool = True,
+            show_line_numbers: bool = True
         ) -> str:
             """Handle diff tool calls."""
             return await tool_self.run(