341 lines
9.2 KiB
Python
341 lines
9.2 KiB
Python
"""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
|
|
) |