code-finder 0.1.0__py3-none-any.whl

claude_context/ast_chunker.py

@@ -0,0 +1,646 @@
+"""
+AST-based Code Chunker for Claude Context
+
+Adapted from the original vibe2doc AST chunker to provide semantic code chunking
+for the Claude Context MCP integration. Uses tree-sitter for accurate parsing
+with a simple line-based fallback for edge cases.
+
+Supports: Python, JavaScript, TypeScript, Go
+"""
+
+import re
+import logging
+from typing import List, Dict, Any, Optional, Tuple
+from dataclasses import dataclass, field
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+# Try to import tree-sitter dependencies
+try:
+    import tree_sitter_python as tspython
+    import tree_sitter_javascript as tsjavascript  
+    import tree_sitter_typescript as tstypescript
+    import tree_sitter_go as tsgo
+    from tree_sitter import Language, Parser, Node
+    HAS_TREE_SITTER = True
+    logger.info("Tree-sitter available for AST parsing")
+except ImportError as e:
+    HAS_TREE_SITTER = False
+    Node = Any  # Type hint fallback
+    logger.warning(f"Tree-sitter not available, using line-based fallback: {e}")
+
+
+@dataclass
+class CodeChunk:
+    """Represents a semantically complete code chunk"""
+    content: str
+    chunk_type: str  # 'function', 'class', 'method', 'import_block', 'module', 'block'
+    name: Optional[str] = None
+    start_line: int = 0
+    end_line: int = 0
+    language: str = ""
+    parent_context: Optional[str] = None
+    docstring: Optional[str] = None  # Extracted docstring content (for functions/classes)
+    size_chars: int = field(default=0, init=False)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    
+    def __post_init__(self):
+        # Count non-whitespace characters for size
+        if self.size_chars == 0:
+            self.size_chars = len(re.sub(r'\s', '', self.content))
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for storage"""
+        return {
+            "content": self.content,
+            "chunk_type": self.chunk_type,
+            "name": self.name,
+            "start_line": self.start_line,
+            "end_line": self.end_line,
+            "language": self.language,
+            "parent_context": self.parent_context,
+            "docstring": self.docstring,
+            "size_chars": self.size_chars,
+            "metadata": self.metadata
+        }
+
+
+class ASTChunker:
+    """
+    AST-based code chunker for semantic code understanding.
+    
+    Primary: Uses tree-sitter for accurate AST parsing
+    Fallback: Simple line-based chunking for edge cases
+    """
+    
+    def __init__(self, max_chunk_size: int = 1500, chunk_overlap: int = 100):
+        """
+        Initialize the AST chunker.
+        
+        Args:
+            max_chunk_size: Maximum chunk size in lines (for fallback)
+            chunk_overlap: Number of lines to overlap (for fallback)
+        """
+        self.max_chunk_size = max_chunk_size
+        self.chunk_overlap = chunk_overlap
+        self.supported_languages = ['python', 'javascript', 'typescript', 'go']
+        self.parsers = {}
+        
+        if HAS_TREE_SITTER:
+            self._initialize_parsers()
+        else:
+            logger.warning("Using line-based chunking fallback")
+    
+    def _initialize_parsers(self):
+        """Initialize tree-sitter parsers for supported languages"""
+        try:
+            # Python parser
+            PY_LANGUAGE = Language(tspython.language())
+            self.parsers['python'] = Parser(PY_LANGUAGE)
+            logger.debug("Python parser initialized")
+        except Exception as e:
+            logger.warning(f"Failed to initialize Python parser: {e}")
+        
+        try:
+            # JavaScript parser  
+            JS_LANGUAGE = Language(tsjavascript.language())
+            self.parsers['javascript'] = Parser(JS_LANGUAGE)
+            logger.debug("JavaScript parser initialized")
+        except Exception as e:
+            logger.warning(f"Failed to initialize JavaScript parser: {e}")
+        
+        try:
+            # TypeScript parser (uses tsx variant)
+            TS_LANGUAGE = Language(tstypescript.language_tsx())
+            self.parsers['typescript'] = Parser(TS_LANGUAGE)
+            logger.debug("TypeScript parser initialized")
+        except Exception as e:
+            logger.warning(f"Failed to initialize TypeScript parser: {e}")
+        
+        try:
+            # Go parser
+            GO_LANGUAGE = Language(tsgo.language())
+            self.parsers['go'] = Parser(GO_LANGUAGE)
+            logger.debug("Go parser initialized")
+        except Exception as e:
+            logger.warning(f"Failed to initialize Go parser: {e}")
+    
+    def chunk_file(self, file_path: Path, content: Optional[str] = None) -> List[CodeChunk]:
+        """
+        Chunk a file into semantic units.
+        
+        Args:
+            file_path: Path to the file
+            content: Optional content (if already loaded)
+            
+        Returns:
+            List of CodeChunk objects
+        """
+        if content is None:
+            try:
+                content = file_path.read_text(encoding='utf-8', errors='ignore')
+            except Exception as e:
+                logger.error(f"Failed to read {file_path}: {e}")
+                return []
+        
+        # Detect language from extension
+        language = self._detect_language(file_path)
+        
+        # Try AST parsing first
+        if HAS_TREE_SITTER and language in self.parsers:
+            try:
+                return self._chunk_with_ast(content, language, str(file_path))
+            except Exception as e:
+                logger.warning(f"AST parsing failed for {file_path}, using fallback: {e}")
+        
+        # Fallback to line-based chunking
+        return self._chunk_with_lines(content, language, str(file_path))
+    
+    def _detect_language(self, file_path: Path) -> str:
+        """Detect programming language from file extension"""
+        ext_to_lang = {
+            '.py': 'python',
+            '.js': 'javascript', 
+            '.jsx': 'javascript',
+            '.ts': 'typescript',
+            '.tsx': 'typescript',
+            '.go': 'go',
+        }
+        return ext_to_lang.get(file_path.suffix.lower(), 'unknown')
+    
+    def _chunk_with_ast(self, content: str, language: str, file_path: str) -> List[CodeChunk]:
+        """
+        Chunk code using tree-sitter AST parsing.
+
+        This provides semantic chunking that respects:
+        - Module-level docstrings (extracted first for Python)
+        - Function boundaries
+        - Class definitions
+        - Method boundaries
+        - Import blocks
+        """
+        if language not in self.parsers:
+            raise ValueError(f"No parser available for {language}")
+
+        parser = self.parsers[language]
+        tree = parser.parse(bytes(content, "utf8"))
+
+        chunks = []
+
+        # Extract module-level docstring for Python files
+        if language == 'python':
+            module_docstring = self._extract_module_docstring(tree.root_node, content)
+            if module_docstring and len(module_docstring.strip()) > 20:
+                # Create a module chunk with the docstring
+                # Extract module name from file path
+                module_name = Path(file_path).stem
+                chunks.append(CodeChunk(
+                    content=module_docstring,
+                    chunk_type='module_docstring',
+                    name=module_name,
+                    start_line=1,
+                    end_line=module_docstring.count('\n') + 1,
+                    language=language,
+                    parent_context=None,
+                    docstring=module_docstring,
+                    metadata={
+                        "node_type": "module",
+                        "has_error": False,
+                        "has_docstring": True,
+                        "is_module_docstring": True
+                    }
+                ))
+
+        # Extract semantic units from the AST
+        chunks.extend(self._extract_chunks_from_node(
+            tree.root_node, content, language, file_path
+        ))
+
+        return chunks
+    
+    def _extract_chunks_from_node(
+        self, 
+        node: Node, 
+        source: str, 
+        language: str,
+        file_path: str,
+        parent_context: Optional[str] = None
+    ) -> List[CodeChunk]:
+        """Recursively extract chunks from AST nodes"""
+        chunks = []
+        
+        # Determine if this node should be a chunk
+        if self._should_chunk_node(node, language):
+            chunk = self._create_chunk_from_node(node, source, language, parent_context)
+            if chunk:
+                chunks.append(chunk)
+                # Update parent context for children
+                if chunk.name:
+                    parent_context = chunk.name
+        
+        # Process children
+        for child in node.children:
+            chunks.extend(self._extract_chunks_from_node(
+                child, source, language, file_path, parent_context
+            ))
+        
+        return chunks
+    
+    def _should_chunk_node(self, node: Node, language: str) -> bool:
+        """Determine if a node should become its own chunk"""
+        chunk_node_types = {
+            'python': [
+                'function_definition', 'class_definition', 
+                'decorated_definition', 'import_from_statement',
+                'import_statement'
+            ],
+            'javascript': [
+                'function_declaration', 'class_declaration',
+                'method_definition', 'arrow_function',
+                'function_expression', 'import_statement',
+                'export_statement'
+            ],
+            'typescript': [
+                'function_declaration', 'class_declaration',
+                'method_definition', 'arrow_function',
+                'interface_declaration', 'type_alias_declaration',
+                'import_statement', 'export_statement'
+            ],
+            'go': [
+                'function_declaration', 'method_declaration',
+                'type_declaration', 'interface_declaration',
+                'import_declaration'
+            ]
+        }
+        
+        node_types = chunk_node_types.get(language, [])
+        return node.type in node_types
+    
+    def _create_chunk_from_node(
+        self,
+        node: Node,
+        source: str,
+        language: str,
+        parent_context: Optional[str] = None
+    ) -> Optional[CodeChunk]:
+        """Create a CodeChunk from an AST node"""
+        # Extract content
+        start_byte = node.start_byte
+        end_byte = node.end_byte
+        content = source[start_byte:end_byte]
+
+        if not content.strip():
+            return None
+
+        # Determine chunk type and name
+        chunk_type, name = self._analyze_node(node, language)
+
+        # Extract docstring for Python functions and classes
+        docstring = None
+        if language == 'python' and node.type in ['function_definition', 'class_definition']:
+            docstring = self._extract_python_docstring(node, source)
+
+        return CodeChunk(
+            content=content,
+            chunk_type=chunk_type,
+            name=name,
+            start_line=node.start_point[0] + 1,
+            end_line=node.end_point[0] + 1,
+            language=language,
+            parent_context=parent_context,
+            docstring=docstring,
+            metadata={
+                "node_type": node.type,
+                "has_error": node.has_error,
+                "has_docstring": docstring is not None
+            }
+        )
+    
+    def _analyze_node(self, node: Node, language: str) -> Tuple[str, Optional[str]]:
+        """Analyze node to determine chunk type and extract name"""
+        node_type = node.type
+        name = None
+        
+        # Language-specific analysis
+        if language == 'python':
+            if node_type == 'function_definition':
+                chunk_type = 'function'
+                name = self._extract_python_function_name(node)
+            elif node_type == 'class_definition':
+                chunk_type = 'class'
+                name = self._extract_python_class_name(node)
+            elif node_type in ['import_statement', 'import_from_statement']:
+                chunk_type = 'import_block'
+            else:
+                chunk_type = 'module'
+                
+        elif language in ['javascript', 'typescript']:
+            if node_type in ['function_declaration', 'function_expression']:
+                chunk_type = 'function'
+                name = self._extract_js_function_name(node)
+            elif node_type == 'class_declaration':
+                chunk_type = 'class'
+                name = self._extract_js_class_name(node)
+            elif node_type == 'arrow_function':
+                chunk_type = 'function'
+                name = 'arrow_function'
+            elif node_type == 'method_definition':
+                chunk_type = 'method'
+                name = self._extract_js_method_name(node)
+            elif node_type in ['import_statement', 'export_statement']:
+                chunk_type = 'import_block'
+            else:
+                chunk_type = 'module'
+                
+        elif language == 'go':
+            if node_type == 'function_declaration':
+                chunk_type = 'function'
+                name = self._extract_go_function_name(node)
+            elif node_type == 'method_declaration':
+                chunk_type = 'method'
+                name = self._extract_go_method_name(node)
+            elif node_type in ['type_declaration', 'interface_declaration']:
+                chunk_type = 'type'
+                name = self._extract_go_type_name(node)
+            elif node_type == 'import_declaration':
+                chunk_type = 'import_block'
+            else:
+                chunk_type = 'module'
+        else:
+            chunk_type = 'unknown'
+        
+        return chunk_type, name
+    
+    def _extract_python_function_name(self, node: Node) -> Optional[str]:
+        """Extract function name from Python AST node"""
+        for child in node.children:
+            if child.type == 'identifier':
+                return child.text.decode('utf-8')
+        return None
+    
+    def _extract_python_class_name(self, node: Node) -> Optional[str]:
+        """Extract class name from Python AST node"""
+        for child in node.children:
+            if child.type == 'identifier':
+                return child.text.decode('utf-8')
+        return None
+    
+    def _extract_js_function_name(self, node: Node) -> Optional[str]:
+        """Extract function name from JavaScript/TypeScript AST node"""
+        for child in node.children:
+            if child.type == 'identifier':
+                return child.text.decode('utf-8')
+        return None
+    
+    def _extract_js_class_name(self, node: Node) -> Optional[str]:
+        """Extract class name from JavaScript/TypeScript AST node"""
+        for child in node.children:
+            if child.type == 'identifier':
+                return child.text.decode('utf-8')
+        return None
+    
+    def _extract_js_method_name(self, node: Node) -> Optional[str]:
+        """Extract method name from JavaScript/TypeScript AST node"""
+        for child in node.children:
+            if child.type == 'property_identifier':
+                return child.text.decode('utf-8')
+        return None
+    
+    def _extract_go_function_name(self, node: Node) -> Optional[str]:
+        """Extract function name from Go AST node"""
+        for child in node.children:
+            if child.type == 'identifier':
+                return child.text.decode('utf-8')
+        return None
+    
+    def _extract_go_method_name(self, node: Node) -> Optional[str]:
+        """Extract method name from Go AST node"""
+        for child in node.children:
+            if child.type == 'field_identifier':
+                return child.text.decode('utf-8')
+        return None
+    
+    def _extract_go_type_name(self, node: Node) -> Optional[str]:
+        """Extract type name from Go AST node"""
+        spec_list = None
+        for child in node.children:
+            if child.type == 'type_spec':
+                spec_list = child
+                break
+
+        if spec_list:
+            for child in spec_list.children:
+                if child.type == 'type_identifier':
+                    return child.text.decode('utf-8')
+        return None
+
+    def _extract_python_docstring(self, node: Node, source: str) -> Optional[str]:
+        """
+        Extract docstring from a Python function/class node.
+
+        Python docstrings are the first expression_statement containing a string
+        inside the block of a function_definition or class_definition.
+        """
+        for child in node.children:
+            if child.type == 'block':
+                for block_child in child.children:
+                    if block_child.type == 'expression_statement':
+                        for expr_child in block_child.children:
+                            if expr_child.type == 'string':
+                                raw = source[expr_child.start_byte:expr_child.end_byte]
+                                return self._clean_docstring(raw)
+                        # First expression_statement wasn't a string, no docstring
+                        return None
+        return None
+
+    def _extract_module_docstring(self, root_node: Node, source: str) -> Optional[str]:
+        """
+        Extract module-level docstring from the root node of a Python file.
+
+        Module docstrings are the first expression_statement containing a string
+        at the top level of the module (direct child of the module/root node).
+        These often contain:
+        - Overall module purpose and description
+        - Usage examples
+        - Algorithm explanations
+        - Paper references (arXiv, DOI, etc.)
+        """
+        for child in root_node.children:
+            if child.type == 'expression_statement':
+                for expr_child in child.children:
+                    if expr_child.type == 'string':
+                        raw = source[expr_child.start_byte:expr_child.end_byte]
+                        return self._clean_docstring(raw)
+                # First expression_statement wasn't a string, no module docstring
+                return None
+            # Skip comments but stop at any other statement type
+            elif child.type not in ['comment']:
+                return None
+        return None
+
+    def _clean_docstring(self, raw: str) -> str:
+        """Clean docstring by removing quotes and normalizing whitespace."""
+        # Remove triple quotes (both styles)
+        if raw.startswith('"""') and raw.endswith('"""'):
+            raw = raw[3:-3]
+        elif raw.startswith("'''") and raw.endswith("'''"):
+            raw = raw[3:-3]
+        # Remove single quotes (less common but valid)
+        elif raw.startswith('"') and raw.endswith('"'):
+            raw = raw[1:-1]
+        elif raw.startswith("'") and raw.endswith("'"):
+            raw = raw[1:-1]
+        return raw.strip()
+
+    def _chunk_with_lines(self, content: str, language: str, file_path: str) -> List[CodeChunk]:
+        """
+        Simple line-based chunking fallback.
+        
+        This is used when tree-sitter is unavailable or fails.
+        Creates overlapping chunks based on line count.
+        """
+        lines = content.split('\n')
+        chunks = []
+        
+        # Simple heuristic: try to detect function/class boundaries
+        i = 0
+        while i < len(lines):
+            # Look for potential start markers
+            chunk_lines = []
+            chunk_start = i
+            
+            # Determine chunk type based on content
+            first_line = lines[i].strip() if i < len(lines) else ""
+            chunk_type = self._detect_chunk_type_from_line(first_line, language)
+            name = self._extract_name_from_line(first_line, language)
+            
+            # Collect lines for this chunk
+            indent_level = len(lines[i]) - len(lines[i].lstrip()) if i < len(lines) else 0
+            
+            # Add lines until we hit size limit or dedent
+            while i < len(lines) and len(chunk_lines) < self.max_chunk_size:
+                line = lines[i]
+                current_indent = len(line) - len(line.lstrip())
+                
+                # Check for end of logical block (dedent)
+                if chunk_lines and line.strip() and current_indent < indent_level:
+                    # Skip if it's just a blank line
+                    if i + 1 < len(lines):
+                        next_indent = len(lines[i + 1]) - len(lines[i + 1].lstrip())
+                        if next_indent >= indent_level:
+                            chunk_lines.append(line)
+                            i += 1
+                            continue
+                    break
+                
+                chunk_lines.append(line)
+                i += 1
+            
+            # Create chunk if we have content
+            if chunk_lines and any(line.strip() for line in chunk_lines):
+                chunks.append(CodeChunk(
+                    content='\n'.join(chunk_lines),
+                    chunk_type=chunk_type if chunk_type else 'block',
+                    name=name,
+                    start_line=chunk_start + 1,
+                    end_line=chunk_start + len(chunk_lines),
+                    language=language,
+                    metadata={"chunking_method": "line-based"}
+                ))
+            
+            # Move back for overlap if chunk is large
+            if len(chunk_lines) > self.chunk_overlap:
+                i -= self.chunk_overlap
+            
+            # Ensure we make progress
+            if i <= chunk_start:
+                i = chunk_start + 1
+        
+        return chunks
+    
+    def _detect_chunk_type_from_line(self, line: str, language: str) -> Optional[str]:
+        """Simple heuristic to detect chunk type from a line"""
+        line = line.strip()
+        
+        if language == 'python':
+            if line.startswith('def '):
+                return 'function'
+            elif line.startswith('class '):
+                return 'class'
+            elif line.startswith(('import ', 'from ')):
+                return 'import_block'
+                
+        elif language in ['javascript', 'typescript']:
+            if 'function ' in line or 'const ' in line and '=>' in line:
+                return 'function'
+            elif line.startswith('class '):
+                return 'class'
+            elif line.startswith(('import ', 'export ')):
+                return 'import_block'
+                
+        elif language == 'go':
+            if line.startswith('func '):
+                return 'function'
+            elif line.startswith('type '):
+                return 'type'
+            elif line.startswith('import'):
+                return 'import_block'
+        
+        return None
+    
+    def _extract_name_from_line(self, line: str, language: str) -> Optional[str]:
+        """Simple heuristic to extract name from a line"""
+        line = line.strip()
+        
+        if language == 'python':
+            if line.startswith('def '):
+                match = re.match(r'def\s+(\w+)', line)
+                return match.group(1) if match else None
+            elif line.startswith('class '):
+                match = re.match(r'class\s+(\w+)', line)
+                return match.group(1) if match else None
+                
+        elif language in ['javascript', 'typescript']:
+            if 'function ' in line:
+                match = re.search(r'function\s+(\w+)', line)
+                return match.group(1) if match else None
+            elif line.startswith('class '):
+                match = re.match(r'class\s+(\w+)', line)
+                return match.group(1) if match else None
+            elif 'const ' in line and '=>' in line:
+                match = re.match(r'const\s+(\w+)', line)
+                return match.group(1) if match else None
+                
+        elif language == 'go':
+            if line.startswith('func '):
+                match = re.match(r'func\s+(?:\([^)]+\)\s+)?(\w+)', line)
+                return match.group(1) if match else None
+            elif line.startswith('type '):
+                match = re.match(r'type\s+(\w+)', line)
+                return match.group(1) if match else None
+        
+        return None
+
+
+# Convenience function for easy import
+def chunk_code(
+    file_path: Path,
+    content: Optional[str] = None,
+    max_chunk_size: int = 1500,
+    chunk_overlap: int = 100
+) -> List[CodeChunk]:
+    """
+    Convenience function to chunk code from a file.
+    
+    Args:
+        file_path: Path to the file
+        content: Optional pre-loaded content
+        max_chunk_size: Maximum lines per chunk (for fallback)
+        chunk_overlap: Line overlap (for fallback)
+        
+    Returns:
+        List of CodeChunk objects
+    """
+    chunker = ASTChunker(max_chunk_size, chunk_overlap)
+    return chunker.chunk_file(file_path, content)