code_extractor.py

import os
import json
import re
from typing import List, Dict, Optional

class CodeExtractor:
    def __init__(self, repo_path: str):
        self.repo_path = repo_path
        self.file_cache = {}

    def load_file_content(self, file_path: str) -> str:
        """Load and cache file content."""
        if file_path not in self.file_cache:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    self.file_cache[file_path] = f.read()
            except Exception as e:
                print(f"Error reading file {file_path}: {e}")
                return ""
        return self.file_cache[file_path]

    def find_code_sections(self, selected_text: str) -> List[Dict]:
        """
        Find code sections containing the selected text.
        Returns a list of dictionaries containing file path and code context.
        """
        results = []
        selected_text = selected_text.strip()
        
        if not selected_text:
            return results

        # Walk through the repository
        for root, _, files in os.walk(self.repo_path):
            for file in files:
                if file.endswith(('.py', '.js', '.jsx', '.ts', '.tsx', '.html', '.css')):
                    file_path = os.path.join(root, file)
                    content = self.load_file_content(file_path)
                    
                    if selected_text in content:
                        # Find the line number
                        lines = content.split('\n')
                        for i, line in enumerate(lines, 1):
                            if selected_text in line:
                                # Get context (5 lines before and after)
                                start = max(0, i - 6)
                                end = min(len(lines), i + 5)
                                context = '\n'.join(lines[start:end])
                                
                                results.append({
                                    'file_path': os.path.relpath(file_path, self.repo_path),
                                    'line_number': i,
                                    'code_context': context,
                                    'selected_text': selected_text
                                })

        return results

    def format_results(self, results: List[Dict]) -> str:
        """Format the results into a readable string."""
        if not results:
            return "No code sections found containing the selected text."

        formatted_output = []
        for result in results:
            formatted_output.append(f"File: {result['file_path']}")
            formatted_output.append(f"Line: {result['line_number']}")
            formatted_output.append("Code Context:")
            formatted_output.append("-" * 80)
            formatted_output.append(result['code_context'])
            formatted_output.append("-" * 80)
            formatted_output.append("")

        return "\n".join(formatted_output)

def extract_code_from_text(repo_path: str, selected_text: str) -> str:
    """
    Main function to extract code sections containing the selected text.
    
    Args:
        repo_path: Path to the repository
        selected_text: Text selected by the user
        
    Returns:
        Formatted string containing the code sections
    """
    extractor = CodeExtractor(repo_path)
    results = extractor.find_code_sections(selected_text)
    return extractor.format_results(results)