scry-run 0.1.0__py3-none-any.whl

scry_run/context.py

@@ -0,0 +1,298 @@
+"""Context builder for collecting codebase and inserting hole markers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+
+class ContextBuilder:
+    """Builds context strings from the codebase for LLM prompts.
+    
+    Collects all Python files in a project and formats them as a single
+    context string, with a "hole" marker indicating where code is missing.
+    """
+    
+    # Directories to always skip
+    SKIP_DIRS = {
+        "__pycache__",
+        ".git",
+        ".venv",
+        "venv",
+        ".env",
+        "node_modules",
+        ".pytest_cache",
+        ".mypy_cache",
+        ".ruff_cache",
+        "dist",
+        "build",
+        "*.egg-info",
+        ".scry-run-cache",
+    }
+    
+    # File patterns to skip
+    SKIP_FILES = {
+        "*.pyc",
+        "*.pyo",
+        "*.pyd",
+    }
+    
+    def __init__(self, root_dir: str | Path | None = None):
+        """Initialize context builder.
+        
+        Args:
+            root_dir: Root directory of the project. Defaults to current directory.
+        """
+        self.root_dir = Path(root_dir) if root_dir else Path.cwd()
+    
+    def _should_skip_dir(self, dir_path: Path) -> bool:
+        """Check if a directory should be skipped."""
+        name = dir_path.name
+        return name in self.SKIP_DIRS or name.startswith(".")
+    
+    def _should_skip_file(self, file_path: Path) -> bool:
+        """Check if a file should be skipped."""
+        name = file_path.name
+        return any(
+            name.endswith(pattern.lstrip("*"))
+            for pattern in self.SKIP_FILES
+        )
+    
+    def collect_python_files(self) -> list[Path]:
+        """Collect all Python files in the project.
+        
+        Returns:
+            List of paths to Python files, relative to root_dir
+        """
+        files = []
+        
+        def walk(directory: Path) -> None:
+            try:
+                for item in sorted(directory.iterdir()):
+                    if item.is_dir():
+                        if not self._should_skip_dir(item):
+                            walk(item)
+                    elif item.is_file() and item.suffix == ".py":
+                        if not self._should_skip_file(item):
+                            files.append(item)
+            except PermissionError:
+                pass
+        
+        walk(self.root_dir)
+        return files
+    
+    def read_file_content(self, file_path: Path) -> str:
+        """Read content of a file.
+        
+        Args:
+            file_path: Path to the file
+            
+        Returns:
+            File content as string
+        """
+        try:
+            return file_path.read_text(encoding="utf-8")
+        except (UnicodeDecodeError, PermissionError, FileNotFoundError):
+            return ""
+    
+    def build_context(
+        self,
+        class_name: str,
+        attr_name: str,
+        class_source: Optional[str] = None,
+    ) -> str:
+        """Build full context with hole marker.
+        
+        Args:
+            class_name: Name of the class needing the attribute
+            attr_name: Name of the missing attribute
+            class_source: Optional source code of the class itself
+            
+        Returns:
+            Formatted context string with all files and hole marker
+        """
+        files = self.collect_python_files()
+        
+        sections = []
+        
+        # Add header
+        sections.append(f"# Project: {self.root_dir.name}")
+        sections.append("")
+        
+        # Add each file
+        for file_path in files:
+            relative_path = file_path.relative_to(self.root_dir)
+            content = self.read_file_content(file_path)
+            
+            if content:
+                sections.append(f"# --- {relative_path} ---")
+                sections.append(content)
+                sections.append("")
+        
+        # Add hole marker
+        sections.append("# --- MISSING CODE ---")
+        sections.append(f"# Class: {class_name}")
+        sections.append(f"# Missing attribute: {attr_name}")
+        sections.append("#")
+        sections.append("# Generate the code for this attribute based on:")
+        sections.append("# 1. The class name and its docstring")
+        sections.append("# 2. Existing methods in the class")
+        sections.append("# 3. How the class is used in the codebase")
+        sections.append("# 4. The attribute name (infer purpose from naming)")
+        
+        if class_source:
+            sections.append("")
+            sections.append("# Current class definition:")
+            sections.append(class_source)
+        
+        sections.append("# --- END MISSING CODE ---")
+        
+        return "\n".join(sections)
+    
+    def build_minimal_context(
+        self,
+        class_source: str,
+        class_name: str,
+        attr_name: str,
+        additional_context: Optional[str] = None,
+    ) -> str:
+        """Build minimal context with just the class and hole marker.
+        
+        Useful for smaller, faster generation when full codebase isn't needed.
+        
+        Args:
+            class_source: Source code of the class
+            class_name: Name of the class
+            attr_name: Name of the missing attribute
+            additional_context: Optional additional context to include
+            
+        Returns:
+            Formatted context string
+        """
+        sections = [
+            "# Class definition:",
+            class_source,
+            "",
+            "# --- MISSING CODE ---",
+            f"# Class: {class_name}",
+            f"# Missing attribute: {attr_name}",
+        ]
+        
+        if additional_context:
+            sections.extend([
+                "",
+                "# Additional context:",
+                additional_context,
+            ])
+        
+        sections.append("# --- END MISSING CODE ---")
+        
+        return "\n".join(sections)
+
+    @staticmethod
+    def build_runtime_context() -> str:
+        """Capture runtime context from the call stack.
+        
+        Captures the call stack and local variables to give the LLM
+        context about how the missing method is being called.
+        
+        Returns:
+            Formatted string with call stack and variable info
+        """
+        import sys
+        import traceback
+        
+        lines = ["# === RUNTIME CONTEXT ===", ""]
+        
+        # Get the current frame and walk up the stack
+        frame = sys._getframe()
+        
+        # Collect relevant frames (skip scry_run internals)
+        relevant_frames = []
+        current = frame
+        while current and len(relevant_frames) < 50:
+            filename = current.f_code.co_filename
+            # Skip scry_run internal frames (but include test files)
+            if "scry_run" not in filename or "test" in filename:
+                relevant_frames.append(current)
+            current = current.f_back
+        
+        if relevant_frames:
+            lines.append("# Call stack (most recent call first):")
+            lines.append("")
+            
+            for i, f in enumerate(relevant_frames[:50]):  # Limit to 50 frames
+                filename = f.f_code.co_filename
+                lineno = f.f_lineno
+                funcname = f.f_code.co_name
+                
+                # Try to get the source line
+                try:
+                    import linecache
+                    source_line = linecache.getline(filename, lineno).strip()
+                except Exception:
+                    source_line = "<source unavailable>"
+                
+                lines.append(f"# Frame {i}: {funcname}() at {filename}:{lineno}")
+                lines.append(f"#   {source_line}")
+                
+                # Capture local variables (but not too many)
+                locals_dict = f.f_locals
+                if locals_dict:
+                    # Filter out internal/large objects
+                    interesting_vars = {}
+                    for name, value in locals_dict.items():
+                        if name.startswith("_"):
+                            continue
+                        # Skip self/cls as they're obvious
+                        if name in ("self", "cls"):
+                            continue
+                        # Try to get a reasonable repr
+                        try:
+                            value_repr = repr(value)
+                            if len(value_repr) > 200:
+                                value_repr = value_repr[:200] + "..."
+                            interesting_vars[name] = value_repr
+                        except Exception:
+                            interesting_vars[name] = "<unable to repr>"
+                    
+                    if interesting_vars:
+                        lines.append(f"#   Local variables:")
+                        for name, val in list(interesting_vars.items())[:5]:
+                            lines.append(f"#     {name} = {val}")
+                
+                lines.append("")
+        
+        # Add the immediate caller's arguments if this is a method call
+        if relevant_frames:
+            caller = relevant_frames[0]
+            caller_locals = caller.f_locals
+            
+            # Get function arguments from the frame
+            code = caller.f_code
+            arg_names = code.co_varnames[:code.co_argcount]
+            
+            args_info = {}
+            for arg_name in arg_names:
+                if arg_name in ("self", "cls"):
+                    continue
+                if arg_name in caller_locals:
+                    try:
+                        val = repr(caller_locals[arg_name])
+                        if len(val) > 200:
+                            val = val[:200] + "..."
+                        args_info[arg_name] = val
+                    except Exception:
+                        pass
+            
+            if args_info:
+                lines.append("# === CALL ARGUMENTS ===")
+                lines.append("# The method was called with these arguments:")
+                for name, val in args_info.items():
+                    lines.append(f"#   {name} = {val}")
+                lines.append("")
+        
+        lines.append("# === END RUNTIME CONTEXT ===")
+        
+        return "\n".join(lines)
+