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

jarvis/jarvis_treesitter/README.md

@@ -0,0 +1,104 @@
+# Tree-sitter 代码数据库
+
+基于 tree-sitter 的代码分析工具，支持快速查询符号的定义位置、声明位置、引用位置和调用关系。
+
+## 功能特点
+
+- 支持多种编程语言：Python、C、C++、Go、Rust
+- 自动下载和编译语言语法文件
+- 查找符号定义
+- 查找符号引用
+- 查找函数调用者
+
+## 安装
+
+```bash
+pip install -r requirements.txt
+```
+
+## 使用方法
+
+### 基本用法
+
+```python
+from jarvis.jarvis_treesitter import CodeDatabase
+
+# 初始化代码数据库（自动下载所需的语法文件）
+db = CodeDatabase()  # 语法文件将保存到 ~/.jarvis/treesitter 目录
+
+# 索引源文件
+db.index_file("path/to/file.py")
+
+# 查找符号
+symbols = db.find_symbol("function_name")
+
+# 查找符号引用
+references = db.find_references(symbols[0])
+
+# 查找函数调用者
+callers = db.find_callers(symbols[0])
+```
+
+### 自定义语法文件位置
+
+虽然默认会使用 `~/.jarvis/treesitter` 目录，但您仍然可以指定自定义目录：
+
+```python
+from jarvis.jarvis_treesitter import CodeDatabase
+
+# 使用自定义语法文件目录
+db = CodeDatabase(grammar_dir="/path/to/grammars")
+
+# 不自动下载缺失的语法文件
+db = CodeDatabase(auto_download=False)
+```
+
+### 手动下载语法文件
+
+```python
+from jarvis.jarvis_treesitter import setup_default_grammars, GrammarBuilder, LanguageType, DEFAULT_GRAMMAR_DIR
+
+# 下载所有支持的语言的语法文件到默认目录 (~/.jarvis/treesitter)
+setup_default_grammars()
+
+# 或者使用自定义目录
+grammar_dir = "/path/to/grammars"
+builder = GrammarBuilder(grammar_dir)
+builder.ensure_all_grammars()  # 下载所有语言
+builder.ensure_grammar(LanguageType.PYTHON)  # 只下载特定语言
+
+# 查看默认语法文件目录
+print(DEFAULT_GRAMMAR_DIR)  # 输出: ~/.jarvis/treesitter
+```
+
+## 命令行工具
+
+提供了一个示例脚本 `example.py` 演示基本用法：
+
+```bash
+# 索引当前目录并查找名为 "main" 的符号
+python -m jarvis.jarvis_treesitter.example --dir . --symbol main
+
+# 只索引Python文件
+python -m jarvis.jarvis_treesitter.example --dir . --ext .py --symbol main
+
+# 使用自定义语法文件目录（默认是 ~/.jarvis/treesitter）
+python -m jarvis.jarvis_treesitter.example --dir . --grammar-dir /path/to/grammars --symbol main
+
+# 不自动下载语法文件
+python -m jarvis.jarvis_treesitter.example --dir . --no-download --symbol main
+```
+
+## 语法文件位置
+
+默认情况下，所有tree-sitter语法文件将保存在 `~/.jarvis/treesitter` 目录中。这些文件只需要下载和编译一次，后续使用时会自动加载。
+
+## 支持的语言
+
+| 语言   | 文件扩展名                | 支持的符号类型                                   |
+|--------|--------------------------|--------------------------------------------------|
+| Python | .py                      | 函数、类、变量、导入、方法                        |
+| C      | .c, .h                   | 函数、结构体、枚举、类型定义、宏、变量            |
+| C++    | .cpp, .hpp, .cc, .hh     | 函数、类、结构体、枚举、命名空间、模板、变量      |
+| Go     | .go                      | 函数、结构体、接口、包、导入、变量                |
+| Rust   | .rs                      | 函数、结构体、枚举、特征、实现、模块、变量        | 

jarvis/jarvis_treesitter/__init__.py

@@ -0,0 +1,20 @@
+"""Tree-sitter based code database for fast symbol lookup."""
+
+__version__ = "0.1.0"
+
+from .database import CodeDatabase
+from .symbol import Symbol, SymbolType, SymbolLocation
+from .language import LanguageType, LanguageConfig
+from .grammar_builder import GrammarBuilder, setup_default_grammars, DEFAULT_GRAMMAR_DIR
+
+__all__ = [
+    "CodeDatabase",
+    "Symbol",
+    "SymbolType",
+    "SymbolLocation",
+    "LanguageType",
+    "LanguageConfig",
+    "GrammarBuilder",
+    "setup_default_grammars",
+    "DEFAULT_GRAMMAR_DIR",
+] 

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