diff --git a/local_code_assistant/commands/refactor.py b/local_code_assistant/commands/refactor.py new file mode 100644 index 0000000..c58edb5 --- /dev/null +++ b/local_code_assistant/commands/refactor.py @@ -0,0 +1,341 @@ +"""Refactor command for Local Code Assistant.""" + +from pathlib import Path +from typing import Optional + +import click +from rich.console import Console +from rich.panel import Panel +from rich.syntax import Syntax + +from local_code_assistant.commands.base import BaseCommand +from local_code_assistant.prompts.templates import LanguageConfig, PromptTemplates +from local_code_assistant.services.ollama import OllamaService + +console = Console() + + +class RefactorCommand(BaseCommand): + """Command for refactoring code.""" + + def __init__(self, ollama: OllamaService, config): + """Initialize refactor command. + + Args: + ollama: Ollama service instance. + config: Configuration service instance. + """ + super().__init__(ollama, config) + self.supported_languages = LanguageConfig.get_supported_languages() + + def run( + self, + code: str, + language: str, + focus: list[str], + output: Optional[Path] = None, + clipboard: bool = False, + safe: bool = True + ) -> str: + """Execute code refactoring. + + Args: + code: Code to refactor. + language: Programming language. + focus: Areas to focus on. + output: Output file path. + clipboard: Copy to clipboard. + safe: Perform safe refactoring only. + + Returns: + Refactored code. + """ + if not self.ollama.check_connection(): + raise click.ClickException( + "Cannot connect to Ollama. Make sure it's running." + ) + + if language not in self.supported_languages: + raise click.ClickException(f"Unsupported language: {language}") + + console.print("[dim]Refactoring code...[/dim]") + + if safe: + full_prompt = PromptTemplates.code_refactor_safe( + language=language, + code=code + ) + else: + full_prompt = PromptTemplates.code_refactor( + language=language, + code=code, + focus=focus + ) + + system_prompt = PromptTemplates.build_system_prompt( + "You are refactoring code. Maintain functionality while improving structure." + ) + + try: + refactored_code = self.ollama.generate( + prompt=full_prompt, + model=self._get_model(), + system=system_prompt, + temperature=self._get_temperature(0.3) + ) + + self._display_output(refactored_code, language) + + if output: + output.write_text(refactored_code) + console.print(f"[green]Refactored code written to {output}[/green]") + + if clipboard: + import pyperclip + pyperclip.copy(refactored_code) + console.print("[green]Refactored code copied to clipboard[/green]") + + return refactored_code + + except Exception as e: + raise click.ClickException(f"Refactoring failed: {str(e)}") from e + + def run_file( + self, + file: Path, + focus: list[str], + output: Optional[Path] = None, + clipboard: bool = False, + safe: bool = True + ) -> str: + """Refactor code from a file. + + Args: + file: Path to file. + focus: Areas to focus on. + output: Output file path. + clipboard: Copy to clipboard. + safe: Perform safe refactoring only. + + Returns: + Refactored code. + """ + if not file.exists(): + raise click.ClickException(f"File not found: {file}") + + code = file.read_text() + language = self._detect_language_from_file(file) + + return self.run( + code=code, + language=language, + focus=focus, + output=output, + clipboard=clipboard, + safe=safe + ) + + def run_optimize( + self, + code: str, + language: str, + output: Optional[Path] = None, + clipboard: bool = False + ) -> str: + """Optimize code for performance. + + Args: + code: Code to optimize. + language: Programming language. + output: Output file path. + clipboard: Copy to clipboard. + + Returns: + Optimized code. + """ + if not self.ollama.check_connection(): + raise click.ClickException( + "Cannot connect to Ollama. Make sure it's running." + ) + + console.print("[dim]Optimizing code...[/dim]") + + full_prompt = PromptTemplates.code_refactor_optimize( + language=language, + code=code + ) + + system_prompt = PromptTemplates.build_system_prompt( + "You are optimizing code for performance. Focus on algorithmic improvements." + ) + + try: + optimized_code = self.ollama.generate( + prompt=full_prompt, + model=self._get_model(), + system=system_prompt, + temperature=self._get_temperature(0.2) + ) + + self._display_output(optimized_code, language) + + if output: + output.write_text(optimized_code) + console.print(f"[green]Optimized code written to {output}[/green]") + + if clipboard: + import pyperclip + pyperclip.copy(optimized_code) + console.print("[green]Optimized code copied to clipboard[/green]") + + return optimized_code + + except Exception as e: + raise click.ClickException(f"Optimization failed: {str(e)}") from e + + def _detect_language_from_file(self, file: Path) -> str: + """Detect language from file extension. + + Args: + file: File path. + + Returns: + Language name. + """ + suffix = file.suffix.lower() + + language_map = { + ".py": "python", + ".js": "javascript", + ".ts": "typescript", + ".tsx": "typescript", + ".go": "go", + ".rs": "rust" + } + + return language_map.get(suffix, "python") + + def _display_output(self, code: str, language: str): + """Display refactored code. + + Args: + code: Refactored code. + language: Programming language. + """ + if self.config.syntax_highlighting: + syntax = Syntax(code, language, line_numbers=True) + console.print(Panel(syntax, title="Refactored Code")) + else: + console.print(Panel(code, title="Refactored Code")) + + +@click.command() +@click.argument("file", type=click.Path(exists=True, path_type=Path)) +@click.option( + "--focus", "-f", + multiple=True, + default=["readability"], + help="Focus areas: readability, structure, naming, documentation" +) +@click.option( + "--safe/--unsafe", + default=True, + help="Safe refactoring (maintains behavior) vs full refactoring" +) +@click.option( + "--output", "-o", + type=click.Path(path_type=Path), + help="Write refactored code to file" +) +@click.option( + "--clipboard/--no-clipboard", + default=False, + help="Copy refactored code to clipboard" +) +@click.pass_context +def refactor_cmd( + ctx: click.Context, + file: Path, + focus: tuple, + safe: bool, + output: Optional[Path], + clipboard: bool +): + """Refactor code to improve structure and readability. + + Example: + local-code-assistant refactor script.py --safe + local-code-assistant refactor app.py -f readability -f naming -o refactored.py + + \f + Args: + ctx: Click context. + file: Path to file to refactor. + focus: Focus areas. + safe: Perform safe refactoring. + output: Output file path. + clipboard: Copy to clipboard. + """ + config = ctx.obj["config"] + ollama_service = OllamaService(config) + ctx.obj["ollama_service"] = ollama_service + + command = RefactorCommand(ollama_service, config) + command.run_file( + file=file, + focus=list(focus), + output=output, + clipboard=clipboard, + safe=safe + ) + + +@click.command() +@click.argument("file", type=click.Path(exists=True, path_type=Path)) +@click.option( + "--output", "-o", + type=click.Path(path_type=Path), + help="Write optimized code to file" +) +@click.option( + "--clipboard/--no-clipboard", + default=False, + help="Copy optimized code to clipboard" +) +@click.pass_context +def optimize_cmd( + ctx: click.Context, + file: Path, + output: Optional[Path], + clipboard: bool +): + """Optimize code for better performance. + + Example: + local-code-assistant optimize slow.py -o fast.py + + \f + Args: + ctx: Click context. + file: Path to file to optimize. + output: Output file path. + clipboard: Copy to clipboard. + """ + config = ctx.obj["config"] + ollama_service = OllamaService(config) + ctx.obj["ollama_service"] = ollama_service + + command = RefactorCommand(ollama_service, config) + + if not file.exists(): + raise click.ClickException(f"File not found: {file}") + + code = file.read_text() + language = command._detect_language_from_file(file) + + command.run_optimize( + code=code, + language=language, + output=output, + clipboard=clipboard + ) \ No newline at end of file