jarvis-ai-assistant 0.1.130__py3-none-any.whl → 0.1.131__py3-none-any.whl

jarvis/jarvis_treesitter/database.py

@@ -0,0 +1,258 @@
+"""Tree-sitter based code database implementation."""
+
+import os
+import logging
+from typing import Dict, List, Optional, Set, Tuple
+from tree_sitter import Language, Node, Parser, Tree
+from .symbol import Symbol, SymbolType, SymbolLocation
+from .language import LanguageType, detect_language, get_language_config
+from .grammar_builder import GrammarBuilder, setup_default_grammars, DEFAULT_GRAMMAR_DIR
+
+# Setup logging
+logger = logging.getLogger(__name__)
+
+class CodeDatabase:
+    """A database for storing and querying code symbols using tree-sitter."""
+    
+    def __init__(self, grammar_dir: Optional[str] = None, auto_download: bool = True):
+        """Initialize the code database.
+        
+        Args:
+            grammar_dir: Directory containing tree-sitter grammar files.
+                         If None, uses the default directory (~/.jarvis/treesitter).
+            auto_download: Whether to automatically download missing grammar files.
+        """
+        self.parser = Parser()
+        self.languages: Dict[LanguageType, Language] = {}
+        self.file_languages: Dict[str, LanguageType] = {}
+        
+        # Use default grammar directory if not provided
+        if grammar_dir is None:
+            grammar_dir = DEFAULT_GRAMMAR_DIR
+            if auto_download:
+                grammar_dir = setup_default_grammars()
+        
+        # Create grammar builder
+        self.grammar_builder = GrammarBuilder(grammar_dir)
+        
+        # Load all supported language grammars
+        for lang_type in LanguageType:
+            try:
+                if auto_download:
+                    # Ensure grammar exists (download if needed)
+                    grammar_path = self.grammar_builder.ensure_grammar(lang_type)
+                else:
+                    # Just check if grammar file exists
+                    config = get_language_config(lang_type)
+                    grammar_path = os.path.join(grammar_dir, config.grammar_file)
+                    if not os.path.exists(grammar_path):
+                        logger.warning(f"Grammar file for {lang_type.value} not found: {grammar_path}")
+                        continue
+                
+                # Load the language
+                config = get_language_config(lang_type)
+                self.languages[lang_type] = Language(grammar_path, config.name)
+                logger.info(f"Loaded language grammar for {lang_type.value}")
+            except Exception as e:
+                logger.error(f"Failed to load grammar for {lang_type.value}: {str(e)}")
+        
+        # Symbol storage
+        self.symbols: Dict[str, List[Symbol]] = {}
+        self.file_trees: Dict[str, Tuple[Tree, LanguageType]] = {}
+        
+    def index_file(self, file_path: str) -> None:
+        """Index a source code file.
+        
+        Args:
+            file_path: Path to the source code file
+        """
+        # Detect language
+        lang_type = detect_language(file_path)
+        if not lang_type:
+            raise ValueError(f"Could not detect language for file: {file_path}")
+            
+        # Check if language is supported
+        if lang_type not in self.languages:
+            if not hasattr(self, 'grammar_builder'):
+                raise ValueError(f"Unsupported language for file: {file_path}")
+            
+            # Try to build the grammar on-demand
+            try:
+                grammar_path = self.grammar_builder.ensure_grammar(lang_type)
+                config = get_language_config(lang_type)
+                self.languages[lang_type] = Language(grammar_path, config.name)
+                logger.info(f"Built and loaded language grammar for {lang_type.value}")
+            except Exception as e:
+                raise ValueError(f"Failed to build grammar for {lang_type.value}: {str(e)}")
+        
+        # Set language for parsing
+        self.parser.set_language(self.languages[lang_type])
+        self.file_languages[file_path] = lang_type
+        
+        # Parse file
+        with open(file_path, 'rb') as f:
+            source_code = f.read()
+            
+        tree = self.parser.parse(source_code)
+        self.file_trees[file_path] = (tree, lang_type)
+        
+        # Extract symbols from the tree
+        self._extract_symbols(tree, file_path, lang_type)
+        
+    def _extract_symbols(self, tree: Tree, file_path: str, lang_type: LanguageType) -> None:
+        """Extract symbols from a tree-sitter tree.
+        
+        Args:
+            tree: The tree-sitter tree
+            file_path: Path to the source file
+            lang_type: The language type
+        """
+        config = get_language_config(lang_type)
+        
+        def visit_node(node: Node):
+            if not node:
+                return
+                
+            # Extract symbols based on language-specific patterns
+            for symbol_type, patterns in config.symbol_patterns.items():
+                if node.type in patterns:
+                    name_node = None
+                    
+                    # Get the name node based on language-specific rules
+                    if lang_type == LanguageType.PYTHON:
+                        name_node = node.child_by_field_name('name')
+                    elif lang_type in (LanguageType.C, LanguageType.CPP):
+                        if node.type == 'function_definition':
+                            name_node = node.child_by_field_name('declarator')
+                        elif node.type in ('struct_specifier', 'class_specifier'):
+                            name_node = node.child_by_field_name('name')
+                    elif lang_type == LanguageType.GO:
+                        if node.type in ('function_declaration', 'method_declaration'):
+                            name_node = node.child_by_field_name('name')
+                    elif lang_type == LanguageType.RUST:
+                        if node.type in ('function_item', 'struct_item', 'enum_item', 'trait_item'):
+                            name_node = node.child_by_field_name('name')
+                    
+                    if name_node and name_node.type == 'identifier':
+                        symbol = Symbol(
+                            name=name_node.text.decode(),
+                            type=SymbolType(symbol_type),
+                            location=SymbolLocation(
+                                file_path=file_path,
+                                start_line=node.start_point[0] + 1,
+                                start_column=node.start_point[1] + 1,
+                                end_line=node.end_point[0] + 1,
+                                end_column=node.end_point[1] + 1
+                            )
+                        )
+                        self._add_symbol(symbol)
+            
+            # Recursively visit children
+            for child in node.children:
+                visit_node(child)
+        
+        visit_node(tree.root_node)
+    
+    def _add_symbol(self, symbol: Symbol) -> None:
+        """Add a symbol to the database.
+        
+        Args:
+            symbol: The symbol to add
+        """
+        if symbol.name not in self.symbols:
+            self.symbols[symbol.name] = []
+        self.symbols[symbol.name].append(symbol)
+    
+    def find_symbol(self, name: str) -> List[Symbol]:
+        """Find all occurrences of a symbol by name.
+        
+        Args:
+            name: The symbol name to search for
+            
+        Returns:
+            List of matching symbols
+        """
+        return self.symbols.get(name, [])
+    
+    def find_references(self, symbol: Symbol) -> List[Symbol]:
+        """Find all references to a symbol.
+        
+        Args:
+            symbol: The symbol to find references for
+            
+        Returns:
+            List of reference symbols
+        """
+        references = []
+        for file_path, (tree, lang_type) in self.file_trees.items():
+            def find_refs(node: Node):
+                if not node:
+                    return
+                    
+                if node.type == 'identifier' and node.text.decode() == symbol.name:
+                    ref_symbol = Symbol(
+                        name=node.text.decode(),
+                        type=SymbolType.REFERENCE,
+                        location=SymbolLocation(
+                            file_path=file_path,
+                            start_line=node.start_point[0] + 1,
+                            start_column=node.start_point[1] + 1,
+                            end_line=node.end_point[0] + 1,
+                            end_column=node.end_point[1] + 1
+                        )
+                    )
+                    references.append(ref_symbol)
+                
+                for child in node.children:
+                    find_refs(child)
+            
+            find_refs(tree.root_node)
+        
+        return references
+    
+    def find_callers(self, function_symbol: Symbol) -> List[Symbol]:
+        """Find all callers of a function.
+        
+        Args:
+            function_symbol: The function symbol to find callers for
+            
+        Returns:
+            List of caller symbols
+        """
+        callers = []
+        for file_path, (tree, lang_type) in self.file_trees.items():
+            def find_calls(node: Node):
+                if not node:
+                    return
+                    
+                # Language-specific call patterns
+                call_patterns = {
+                    LanguageType.PYTHON: ('call', 'function'),
+                    LanguageType.C: ('call_expression', 'function'),
+                    LanguageType.CPP: ('call_expression', 'function'),
+                    LanguageType.GO: ('call_expression', 'function'),
+                    LanguageType.RUST: ('call_expression', 'function'),
+                }
+                
+                if node.type == call_patterns[lang_type][0]:
+                    func_node = node.child_by_field_name(call_patterns[lang_type][1])
+                    if func_node and func_node.type == 'identifier' and func_node.text.decode() == function_symbol.name:
+                        caller_symbol = Symbol(
+                            name=func_node.text.decode(),
+                            type=SymbolType.FUNCTION_CALL,
+                            location=SymbolLocation(
+                                file_path=file_path,
+                                start_line=node.start_point[0] + 1,
+                                start_column=node.start_point[1] + 1,
+                                end_line=node.end_point[0] + 1,
+                                end_column=node.end_point[1] + 1
+                            )
+                        )
+                        callers.append(caller_symbol)
+                
+                for child in node.children:
+                    find_calls(child)
+            
+            find_calls(tree.root_node)
+        
+        return callers 

jarvis/jarvis_treesitter/example.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""Example script demonstrating the use of the tree-sitter code database."""
+
+import os
+import sys
+import logging
+import argparse
+from pathlib import Path
+from typing import List, Optional, Set
+
+from jarvis.jarvis_treesitter import (
+    CodeDatabase, 
+    SymbolType, 
+    setup_default_grammars,
+    DEFAULT_GRAMMAR_DIR
+)
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+def index_directory(db: CodeDatabase, directory: str, extensions: Optional[Set[str]] = None) -> int:
+    """Index all supported files in a directory.
+    
+    Args:
+        db: The code database
+        directory: Directory to index
+        extensions: Optional set of file extensions to index (e.g., {'.py', '.c'})
+        
+    Returns:
+        Number of files indexed
+    """
+    count = 0
+    for root, _, files in os.walk(directory):
+        for file in files:
+            if extensions and not any(file.endswith(ext) for ext in extensions):
+                continue
+                
+            file_path = os.path.join(root, file)
+            try:
+                db.index_file(file_path)
+                count += 1
+                logger.info(f"Indexed file: {file_path}")
+            except Exception as e:
+                logger.warning(f"Failed to index file {file_path}: {str(e)}")
+    
+    return count
+
+def find_symbol(db: CodeDatabase, symbol_name: str) -> None:
+    """Find and print all occurrences of a symbol.
+    
+    Args:
+        db: The code database
+        symbol_name: Symbol name to search for
+    """
+    symbols = db.find_symbol(symbol_name)
+    
+    if not symbols:
+        print(f"No symbols found with name: {symbol_name}")
+        return
+        
+    print(f"Found {len(symbols)} symbols with name: {symbol_name}")
+    for i, symbol in enumerate(symbols):
+        print(f"\n[{i+1}] {symbol.type.value}: {symbol.name}")
+        print(f"    Location: {symbol.location.file_path}:{symbol.location.start_line}:{symbol.location.start_column}")
+        
+        # Find references for this symbol
+        refs = db.find_references(symbol)
+        print(f"    References: {len(refs)}")
+        for j, ref in enumerate(refs[:5]):  # Show first 5 references
+            print(f"      [{j+1}] {ref.location.file_path}:{ref.location.start_line}:{ref.location.start_column}")
+        
+        if len(refs) > 5:
+            print(f"      ... and {len(refs) - 5} more")
+            
+        # Find callers if it's a function
+        if symbol.type == SymbolType.FUNCTION:
+            callers = db.find_callers(symbol)
+            print(f"    Callers: {len(callers)}")
+            for j, caller in enumerate(callers[:5]):  # Show first 5 callers
+                print(f"      [{j+1}] {caller.location.file_path}:{caller.location.start_line}:{caller.location.start_column}")
+            
+            if len(callers) > 5:
+                print(f"      ... and {len(callers) - 5} more")
+
+def main() -> None:
+    """Main entry point."""
+    parser = argparse.ArgumentParser(description="Tree-sitter code database example")
+    parser.add_argument("--dir", "-d", type=str, default=".", help="Directory to index")
+    parser.add_argument("--ext", "-e", type=str, nargs="*", help="File extensions to index (e.g., .py .c)")
+    parser.add_argument("--symbol", "-s", type=str, help="Symbol name to search for")
+    parser.add_argument("--grammar-dir", "-g", type=str, default=DEFAULT_GRAMMAR_DIR, 
+                       help=f"Directory containing grammar files (default: {DEFAULT_GRAMMAR_DIR})")
+    parser.add_argument("--no-download", action="store_true", help="Don't download missing grammars")
+    args = parser.parse_args()
+    
+    # Create code database
+    db = CodeDatabase(grammar_dir=args.grammar_dir, auto_download=not args.no_download)
+    
+    # Process extensions
+    extensions = set(args.ext) if args.ext else None
+    
+    # Index directory
+    count = index_directory(db, args.dir, extensions)
+    print(f"Indexed {count} files in {args.dir}")
+    
+    # Search for symbol if specified
+    if args.symbol:
+        find_symbol(db, args.symbol)
+
+if __name__ == "__main__":
+    main() 

jarvis/jarvis_treesitter/grammar_builder.py

@@ -0,0 +1,182 @@
+"""Functionality for downloading and building tree-sitter grammars."""
+
+import os
+import subprocess
+import tempfile
+import shutil
+from typing import Dict, List, Optional
+import logging
+from pathlib import Path
+
+from .language import LanguageType, get_language_config
+
+# Setup logging
+logger = logging.getLogger(__name__)
+
+# Default grammar directory
+DEFAULT_GRAMMAR_DIR = os.path.expanduser("~/.jarvis/treesitter")
+
+# Tree-sitter grammar repositories
+GRAMMAR_REPOS = {
+    LanguageType.PYTHON: "https://github.com/tree-sitter/tree-sitter-python",
+    LanguageType.C: "https://github.com/tree-sitter/tree-sitter-c",
+    LanguageType.CPP: "https://github.com/tree-sitter/tree-sitter-cpp",
+    LanguageType.GO: "https://github.com/tree-sitter/tree-sitter-go",
+    LanguageType.RUST: "https://github.com/tree-sitter/tree-sitter-rust",
+}
+
+class GrammarBuilder:
+    """Handles downloading and building tree-sitter grammar files."""
+    
+    def __init__(self, grammar_dir: str = DEFAULT_GRAMMAR_DIR):
+        """Initialize the grammar builder.
+        
+        Args:
+            grammar_dir: Directory to store built grammar files. 
+                       Defaults to ~/.jarvis/treesitter
+        """
+        self.grammar_dir = grammar_dir
+        os.makedirs(grammar_dir, exist_ok=True)
+        
+    def ensure_grammar(self, lang_type: LanguageType) -> str:
+        """Ensure the grammar file for a language exists, downloading and building if necessary.
+        
+        Args:
+            lang_type: The language type
+            
+        Returns:
+            Path to the grammar file
+        """
+        config = get_language_config(lang_type)
+        grammar_path = os.path.join(self.grammar_dir, config.grammar_file)
+        
+        # Check if grammar file already exists
+        if os.path.exists(grammar_path):
+            logger.info(f"Grammar file for {lang_type.value} already exists at {grammar_path}")
+            return grammar_path
+            
+        # Download and build the grammar
+        logger.info(f"Building grammar for {lang_type.value}")
+        return self._build_grammar(lang_type)
+    
+    def ensure_all_grammars(self) -> Dict[LanguageType, str]:
+        """Ensure grammar files for all supported languages exist.
+        
+        Returns:
+            Dictionary mapping language types to grammar file paths
+        """
+        result = {}
+        for lang_type in LanguageType:
+            try:
+                path = self.ensure_grammar(lang_type)
+                result[lang_type] = path
+            except Exception as e:
+                logger.error(f"Failed to build grammar for {lang_type.value}: {str(e)}")
+        
+        return result
+    
+    def _build_grammar(self, lang_type: LanguageType) -> str:
+        """Download and build the grammar for a language.
+        
+        Args:
+            lang_type: The language type
+            
+        Returns:
+            Path to the built grammar file
+            
+        Raises:
+            RuntimeError: If grammar building fails
+        """
+        config = get_language_config(lang_type)
+        repo_url = GRAMMAR_REPOS.get(lang_type)
+        
+        if not repo_url:
+            raise ValueError(f"No repository URL defined for language {lang_type.value}")
+        
+        # Create temporary directory
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # Clone the repository
+            logger.info(f"Cloning {repo_url}")
+            result = subprocess.run(
+                ["git", "clone", "--depth", "1", repo_url, temp_dir],
+                check=False,
+                capture_output=True,
+                text=True
+            )
+            
+            if result.returncode != 0:
+                raise RuntimeError(f"Failed to clone repository {repo_url}: {result.stderr}")
+            
+            # Build the grammar
+            grammar_path = os.path.join(self.grammar_dir, config.grammar_file)
+            
+            # Create build script
+            build_script = self._create_build_script(temp_dir, lang_type.value, grammar_path)
+            
+            # Execute build script
+            logger.info(f"Building grammar for {lang_type.value}")
+            result = subprocess.run(
+                ["python", build_script],
+                check=False,
+                capture_output=True,
+                text=True
+            )
+            
+            if result.returncode != 0:
+                raise RuntimeError(f"Failed to build grammar for {lang_type.value}: {result.stderr}")
+            
+            # Verify file exists
+            if not os.path.exists(grammar_path):
+                raise RuntimeError(f"Grammar file {grammar_path} was not created")
+            
+            logger.info(f"Successfully built grammar for {lang_type.value}: {grammar_path}")
+            return grammar_path
+    
+    def _create_build_script(self, repo_dir: str, lang_name: str, output_path: str) -> str:
+        """Create a Python script to build the grammar.
+        
+        Args:
+            repo_dir: Path to the cloned repository
+            lang_name: Language name
+            output_path: Output path for the built grammar
+            
+        Returns:
+            Path to the build script
+        """
+        script_path = os.path.join(repo_dir, "build_grammar.py")
+        
+        with open(script_path, "w") as f:
+            f.write(f'''
+import os
+from tree_sitter import Language
+
+# Ensure output directory exists
+os.makedirs(os.path.dirname("{output_path}"), exist_ok=True)
+
+# Build the language
+Language.build_library(
+    "{output_path}",
+    [
+        "{repo_dir}"
+    ]
+)
+
+print(f"Built grammar: {output_path}")
+''')
+        
+        return script_path
+
+
+def setup_default_grammars() -> str:
+    """Set up default grammars in ~/.jarvis/treesitter directory.
+    
+    Returns:
+        Path to the grammar directory
+    """
+    grammar_dir = DEFAULT_GRAMMAR_DIR
+    os.makedirs(grammar_dir, exist_ok=True)
+    
+    builder = GrammarBuilder(grammar_dir)
+    builder.ensure_all_grammars()
+    
+    return grammar_dir 

jarvis/jarvis_treesitter/language.py

@@ -0,0 +1,117 @@
+"""Language-specific configurations for tree-sitter."""
+
+from enum import Enum
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass
+
+class LanguageType(Enum):
+    """Supported programming languages."""
+    PYTHON = "python"
+    C = "c"
+    CPP = "cpp"
+    GO = "go"
+    RUST = "rust"
+
+@dataclass
+class LanguageConfig:
+    """Configuration for a specific programming language."""
+    name: str
+    grammar_file: str
+    file_extensions: List[str]
+    symbol_patterns: Dict[str, List[str]]  # Maps symbol types to tree-sitter node types
+
+# Language-specific configurations
+LANGUAGE_CONFIGS: Dict[LanguageType, LanguageConfig] = {
+    LanguageType.PYTHON: LanguageConfig(
+        name="python",
+        grammar_file="python.so",
+        file_extensions=[".py"],
+        symbol_patterns={
+            "function": ["function_definition"],
+            "class": ["class_definition"],
+            "variable": ["assignment"],
+            "import": ["import_statement", "import_from_statement"],
+            "method": ["function_definition"],
+        }
+    ),
+    LanguageType.C: LanguageConfig(
+        name="c",
+        grammar_file="c.so",
+        file_extensions=[".c", ".h"],
+        symbol_patterns={
+            "function": ["function_definition"],
+            "struct": ["struct_specifier"],
+            "enum": ["enum_specifier"],
+            "typedef": ["type_definition"],
+            "macro": ["preproc_def"],
+            "variable": ["declaration"],
+        }
+    ),
+    LanguageType.CPP: LanguageConfig(
+        name="cpp",
+        grammar_file="cpp.so",
+        file_extensions=[".cpp", ".hpp", ".cc", ".hh"],
+        symbol_patterns={
+            "function": ["function_definition", "method_definition"],
+            "class": ["class_specifier"],
+            "struct": ["struct_specifier"],
+            "enum": ["enum_specifier"],
+            "namespace": ["namespace_definition"],
+            "template": ["template_declaration"],
+            "variable": ["declaration"],
+        }
+    ),
+    LanguageType.GO: LanguageConfig(
+        name="go",
+        grammar_file="go.so",
+        file_extensions=[".go"],
+        symbol_patterns={
+            "function": ["function_declaration", "method_declaration"],
+            "struct": ["type_declaration"],
+            "interface": ["type_declaration"],
+            "package": ["package_clause"],
+            "import": ["import_declaration"],
+            "variable": ["var_declaration", "short_var_declaration"],
+        }
+    ),
+    LanguageType.RUST: LanguageConfig(
+        name="rust",
+        grammar_file="rust.so",
+        file_extensions=[".rs"],
+        symbol_patterns={
+            "function": ["function_item"],
+            "struct": ["struct_item"],
+            "enum": ["enum_item"],
+            "trait": ["trait_item"],
+            "impl": ["impl_item"],
+            "module": ["mod_item"],
+            "variable": ["let_declaration", "const_item", "static_item"],
+        }
+    ),
+}
+
+def detect_language(file_path: str) -> Optional[LanguageType]:
+    """Detect the programming language of a file based on its extension.
+    
+    Args:
+        file_path: Path to the source file
+        
+    Returns:
+        The detected language type or None if not supported
+    """
+    extension = file_path.lower().split('.')[-1]
+    for lang_type, config in LANGUAGE_CONFIGS.items():
+        if f".{extension}" in config.file_extensions:
+            return lang_type
+    return None
+
+def get_language_config(lang_type: LanguageType) -> LanguageConfig:
+    """Get the configuration for a specific language.
+    
+    Args:
+        lang_type: The language type
+        
+    Returns:
+        The language configuration
+    """
+    return LANGUAGE_CONFIGS[lang_type]