strix/strix/interface/tool_components/file_edit_renderer.py

from functools import cache
from typing import Any, ClassVar

from pygments.lexers import get_lexer_by_name, get_lexer_for_filename
from pygments.styles import get_style_by_name
from pygments.util import ClassNotFound
from rich.text import Text
from textual.widgets import Static

from .base_renderer import BaseToolRenderer
from .registry import register_tool_renderer


@cache
def _get_style_colors() -> dict[Any, str]:
    style = get_style_by_name("native")
    return {token: f"#{style_def['color']}" for token, style_def in style if style_def["color"]}


def _get_lexer_for_file(path: str) -> Any:
    try:
        return get_lexer_for_filename(path)
    except ClassNotFound:
        return get_lexer_by_name("text")


@register_tool_renderer
class StrReplaceEditorRenderer(BaseToolRenderer):
    tool_name: ClassVar[str] = "str_replace_editor"
    css_classes: ClassVar[list[str]] = ["tool-call", "file-edit-tool"]

    @classmethod
    def _get_token_color(cls, token_type: Any) -> str | None:
        colors = _get_style_colors()
        while token_type:
            if token_type in colors:
                return colors[token_type]
            token_type = token_type.parent
        return None

    @classmethod
    def _highlight_code(cls, code: str, path: str) -> Text:
        lexer = _get_lexer_for_file(path)
        text = Text()

        for token_type, token_value in lexer.get_tokens(code):
            if not token_value:
                continue
            color = cls._get_token_color(token_type)
            text.append(token_value, style=color)

        return text

    @classmethod
    def render(cls, tool_data: dict[str, Any]) -> Static:
        args = tool_data.get("args", {})
        result = tool_data.get("result")

        command = args.get("command", "")
        path = args.get("path", "")
        old_str = args.get("old_str", "")
        new_str = args.get("new_str", "")
        file_text = args.get("file_text", "")

        text = Text()

        icons_and_labels = {
            "view": ("◇ ", "read", "#10b981"),
            "str_replace": ("◇ ", "edit", "#10b981"),
            "create": ("◇ ", "create", "#10b981"),
            "insert": ("◇ ", "insert", "#10b981"),
            "undo_edit": ("◇ ", "undo", "#10b981"),
        }

        icon, label, color = icons_and_labels.get(command, ("◇ ", "file", "#10b981"))
        text.append(icon, style=color)
        text.append(label, style="dim")

        if path:
            path_display = path[-60:] if len(path) > 60 else path
            text.append(" ")
            text.append(path_display, style="dim")

        if command == "str_replace" and (old_str or new_str):
            if old_str:
                highlighted_old = cls._highlight_code(old_str, path)
                for line in highlighted_old.plain.split("\n"):
                    text.append("\n")
                    text.append("-", style="#ef4444")
                    text.append(" ")
                    text.append(line)

            if new_str:
                highlighted_new = cls._highlight_code(new_str, path)
                for line in highlighted_new.plain.split("\n"):
                    text.append("\n")
                    text.append("+", style="#22c55e")
                    text.append(" ")
                    text.append(line)

        elif command == "create" and file_text:
            text.append("\n")
            text.append_text(cls._highlight_code(file_text, path))

        elif command == "insert" and new_str:
            highlighted_new = cls._highlight_code(new_str, path)
            for line in highlighted_new.plain.split("\n"):
                text.append("\n")
                text.append("+", style="#22c55e")
                text.append(" ")
                text.append(line)

        elif isinstance(result, str) and result.strip():
            text.append("\n  ")
            text.append(result.strip(), style="dim")
        elif not (result and isinstance(result, dict) and "content" in result) and not path:
            text.append(" ")
            text.append("Processing...", style="dim")

        css_classes = cls.get_css_classes("completed")
        return Static(text, classes=css_classes)


@register_tool_renderer
class ListFilesRenderer(BaseToolRenderer):
    tool_name: ClassVar[str] = "list_files"
    css_classes: ClassVar[list[str]] = ["tool-call", "file-edit-tool"]

    @classmethod
    def render(cls, tool_data: dict[str, Any]) -> Static:
        args = tool_data.get("args", {})
        path = args.get("path", "")

        text = Text()
        text.append("◇ ", style="#10b981")
        text.append("list", style="dim")
        text.append(" ")

        if path:
            path_display = path[-60:] if len(path) > 60 else path
            text.append(path_display, style="dim")
        else:
            text.append("Current directory", style="dim")

        css_classes = cls.get_css_classes("completed")
        return Static(text, classes=css_classes)


@register_tool_renderer
class SearchFilesRenderer(BaseToolRenderer):
    tool_name: ClassVar[str] = "search_files"
    css_classes: ClassVar[list[str]] = ["tool-call", "file-edit-tool"]

    @classmethod
    def render(cls, tool_data: dict[str, Any]) -> Static:
        args = tool_data.get("args", {})
        path = args.get("path", "")
        regex = args.get("regex", "")

        text = Text()
        text.append("◇ ", style="#a855f7")
        text.append("search", style="dim")
        text.append("  ")

        if path and regex:
            text.append(path, style="dim")
            text.append(" ", style="dim")
            text.append(regex, style="#a855f7")
        elif path:
            text.append(path, style="dim")
        elif regex:
            text.append(regex, style="#a855f7")
        else:
            text.append("...", style="dim")

        css_classes = cls.get_css_classes("completed")
        return Static(text, classes=css_classes)