Remove colored and highlight (#508)

mentat/command/commands/redo.py

@@ -15,7 +15,7 @@ async def apply(self, *args: str) -> None:
 
         errors = await code_file_manager.history.redo()
         if errors:
-            stream.send(errors)
+            stream.send("\n".join(errors), style="error")
         stream.send("Redo complete", style="success")
 
     @override

mentat/command/commands/search.py

@@ -1,6 +1,5 @@
 from typing import List, Set
 
-from termcolor import colored
 from typing_extensions import override
 
 from mentat.command.command import Command, CommandArgument
@@ -51,28 +50,30 @@ async def apply(self, *args: str) -> None:
 
         cumulative_tokens = 0
         for i, (feature, _) in enumerate(results, start=1):
-            prefix = "\n   "
+            stream.send(str(i).ljust(3), end="")
+            prefix = "   "
 
+            # TODO: The file_name was originally bolded; I decided it wasn't worth it to add a way to add
+            # text attributes for this one scenario, but if we ever do add text attributes, make this bold again.
             file_name = feature.rel_path(session_context.cwd)
-            file_name = colored(file_name, "blue", attrs=["bold"])
-            file_name += colored(feature.interval_string(), "light_cyan")
+            stream.send(file_name, color="blue", end="")
+            file_interval = feature.interval_string()
+            stream.send(file_interval, color="light_cyan", end="")
 
             tokens = feature.count_tokens(config.model)
             cumulative_tokens += tokens
-            tokens_str = colored(f"  ({tokens} tokens)", "yellow")
-            file_name += tokens_str
+            tokens_str = f"  ({tokens} tokens)"
+            stream.send(tokens_str, color="yellow")
 
-            name = []
             if feature.name:
                 name = feature.name.split(",")
-                name = [
-                    f"{'└' if i == len(name) - 1 else '├'}─ {colored(n, 'cyan')}"
-                    for i, n in enumerate(name)
-                ]
-
-            message = f"{str(i).ljust(3)}" + prefix.join([file_name] + name + [""])
-            stream.send(message)
-            if i > 1 and i % SEARCH_RESULT_BATCH_SIZE == 0:
+                for j, n in enumerate(name):
+                    stream.send(prefix, end="")
+                    stream.send(f"{'└' if j == len(name) - 1 else '├'}─ ", end="")
+                    stream.send(n, color="cyan")
+            stream.send("")
+
+            if i % SEARCH_RESULT_BATCH_SIZE == 0:
                 # Required to avoid circular imports, but not ideal.
                 from mentat.session_input import collect_user_input
 

mentat/command/commands/undo.py

@@ -15,7 +15,7 @@ async def apply(self, *args: str) -> None:
 
         errors = code_file_manager.history.undo()
         if errors:
-            stream.send(errors)
+            stream.send("\n".join(errors), style="error")
         stream.send("Undo complete", style="success")
 
     @override

mentat/command/commands/undoall.py

@@ -15,7 +15,7 @@ async def apply(self, *args: str) -> None:
 
         errors = code_file_manager.history.undo_all()
         if errors:
-            stream.send(errors)
+            stream.send("\n".join(errors), style="error")
         stream.send("Undos complete", style="success")
 
     @override

mentat/edit_history.py

@@ -1,6 +1,4 @@
-from typing import Optional
-
-from termcolor import colored
+from typing import List
 
 from mentat.errors import HistoryError
 from mentat.parsers.file_edit import FileEdit
@@ -22,9 +20,9 @@ def push_edits(self):
             self.edits.append(self.cur_edits)
             self.cur_edits = list[FileEdit]()
 
-    def undo(self) -> str:
+    def undo(self) -> List[str]:
         if not self.edits:
-            return colored("No edits available to undo", color="light_red")
+            return ["No edits available to undo"]
 
         # Make sure to go top down
         cur_edit = self.edits.pop()
@@ -36,14 +34,14 @@ def undo(self) -> str:
                 cur_file_edit.undo()
                 undone_edit.append(cur_file_edit)
             except HistoryError as e:
-                errors.append(colored(str(e), color="light_red"))
+                errors.append(str(e))
         if undone_edit:
             self.undone_edits.append(undone_edit)
-        return "\n".join(errors)
+        return errors
 
-    async def redo(self) -> Optional[str]:
+    async def redo(self) -> List[str]:
         if not self.undone_edits:
-            return colored("No edits available to redo", color="light_red")
+            return ["No edits available to redo"]
 
         session_context = SESSION_CONTEXT.get()
         code_file_manager = session_context.code_file_manager
@@ -53,14 +51,15 @@ async def redo(self) -> Optional[str]:
         for edit in edits_to_redo:
             edit.display_full_edit(code_file_manager.file_lines[edit.file_path])
         await code_file_manager.write_changes_to_files(edits_to_redo)
+        return []
 
-    def undo_all(self) -> str:
+    def undo_all(self) -> List[str]:
         if not self.edits:
-            return colored("No edits available to undo", color="light_red")
+            return ["No edits available to undo"]
 
         errors = list[str]()
         while self.edits:
             error = self.undo()
             if error:
-                errors.append(error)
-        return "\n".join(errors)
+                errors += error
+        return errors

mentat/parsers/block_parser.py

@@ -11,6 +11,7 @@
 from mentat.parsers.change_display_helper import DisplayInformation, FileActionType
 from mentat.parsers.file_edit import FileEdit, Replacement
 from mentat.parsers.parser import ParsedLLMResponse, Parser
+from mentat.parsers.streaming_printer import FormattedString
 from mentat.prompts.prompts import read_prompt
 from mentat.session_context import SESSION_CONTEXT
 
@@ -192,7 +193,7 @@ def _add_code_block(
         code_block: str,
         display_information: DisplayInformation,
         file_edit: FileEdit,
-    ) -> str:
+    ) -> FormattedString:
         file_edit.replacements.append(
             Replacement(
                 display_information.first_changed_line,

mentat/parsers/change_display_helper.py

@@ -1,14 +1,15 @@
 from enum import Enum
 from pathlib import Path
+from typing import Tuple
 
 import attr
-from pygments import highlight  # pyright: ignore[reportUnknownVariableType]
+from pygments import lex
 from pygments.formatters import TerminalFormatter
 from pygments.lexer import Lexer
 from pygments.lexers import TextLexer, get_lexer_for_filename
 from pygments.util import ClassNotFound
-from termcolor import colored
 
+from mentat.parsers.streaming_printer import FormattedString
 from mentat.session_context import SESSION_CONTEXT
 from mentat.utils import get_relative_path
 
@@ -94,7 +95,7 @@ def _remove_extra_empty_lines(lines: list[str]) -> list[str]:
     return lines[max(start - 1, 0) : end + 2]
 
 
-def _prefixed_lines(line_number_buffer: int, lines: list[str], prefix: str):
+def _prefixed_lines(line_number_buffer: int, lines: list[str], prefix: str) -> str:
     return "\n".join(
         [
             prefix + " " * (line_number_buffer - len(prefix)) + line.strip("\n")
@@ -108,16 +109,18 @@ def _get_code_block(
     line_number_buffer: int,
     prefix: str,
     color: str | None,
-):
+) -> FormattedString:
     lines = _prefixed_lines(line_number_buffer, code_lines, prefix)
-    if lines:
-        return "\n".join(colored(line, color=color) for line in lines.split("\n"))
+    if color is None:
+        return lines
     else:
-        return ""
+        return (lines, {"color": color})
+
 
+def display_full_change(display_information: DisplayInformation, prefix: str = ""):
+    ctx = SESSION_CONTEXT.get()
 
-def get_full_change(display_information: DisplayInformation, prefix: str = ""):
-    to_print = [
+    full_change = [
         get_file_name(display_information),
         (
             change_delimiter
@@ -134,42 +137,42 @@ def get_full_change(display_information: DisplayInformation, prefix: str = ""):
             else ""
         ),
     ]
-    full_change = "\n".join([line for line in to_print if line])
-    prefixed_change = "\n".join(
-        (prefix + line) if line.strip() else line for line in full_change.split("\n")
-    )
-    return prefixed_change
+    for line in full_change:
+        if (isinstance(line, str) and line.strip()) or (
+            isinstance(line, Tuple) and line[0].strip()
+        ):
+            ctx.stream.send(prefix, end="")
+            ctx.stream.send(line)
 
 
 def get_file_name(
     display_information: DisplayInformation,
-):
+) -> FormattedString:
     match display_information.file_action_type:
         case FileActionType.CreateFile:
-            return "\n" + colored(
-                f"{display_information.file_name}*", color="light_green"
-            )
+            return (f"\n{display_information.file_name}*", {"color": "light_green"})
         case FileActionType.DeleteFile:
-            return "\n" + colored(
-                f"Deletion: {display_information.file_name}", color="light_red"
+            return (
+                f"\nDeletion: {display_information.file_name}",
+                {"color": "light_red"},
             )
         case FileActionType.RenameFile:
-            return "\n" + colored(
-                f"Rename: {display_information.file_name} ->"
-                f" {display_information.new_name}",
-                color="yellow",
+            return (
+                (
+                    f"\nRename: {display_information.file_name} ->"
+                    f" {display_information.new_name}"
+                ),
+                {"color": "yellow"},
             )
         case FileActionType.UpdateFile:
-            return "\n" + colored(
-                f"{display_information.file_name}", color="light_blue"
-            )
+            return (f"\n{display_information.file_name}", {"color": "light_blue"})
 
 
 def get_added_lines(
     display_information: DisplayInformation,
     prefix: str = "+",
     color: str | None = "green",
-):
+) -> FormattedString:
     return _get_code_block(
         display_information.added_block,
         display_information.line_number_buffer,
@@ -182,7 +185,7 @@ def get_removed_lines(
     display_information: DisplayInformation,
     prefix: str = "-",
     color: str | None = "red",
-):
+) -> FormattedString:
     return _get_code_block(
         display_information.removed_block,
         display_information.line_number_buffer,
@@ -191,15 +194,20 @@ def get_removed_lines(
     )
 
 
-def highlight_text(text: str, lexer: Lexer) -> str:
-    # pygments doesn't have type hints on TerminalFormatter
-    return highlight(text, lexer, TerminalFormatter(bg="dark"))  # type: ignore
+def highlight_text(text: str, lexer: Lexer) -> FormattedString:
+    formatter = TerminalFormatter(bg="dark")  # type: ignore
+    string: FormattedString = []
+    for ttype, value in lex(text, lexer):
+        # We use TerminalFormatter's color scheme; TODO: Hook this up to our style themes instead
+        color = formatter._get_color(ttype)  # type: ignore
+        string.append((value, {"color": color}))
+    return string
 
 
 def get_previous_lines(
     display_information: DisplayInformation,
     num: int = 2,
-) -> str:
+) -> FormattedString:
     if display_information.first_changed_line < 0:
         return ""
     lines = _remove_extra_empty_lines(
@@ -229,7 +237,7 @@ def get_previous_lines(
 def get_later_lines(
     display_information: DisplayInformation,
     num: int = 2,
-) -> str:
+) -> FormattedString:
     if display_information.last_changed_line < 0:
         return ""
     lines = _remove_extra_empty_lines(

mentat/parsers/file_edit.py

@@ -10,7 +10,7 @@
     DisplayInformation,
     FileActionType,
     change_delimiter,
-    get_full_change,
+    display_full_change,
 )
 from mentat.session_context import SESSION_CONTEXT
 from mentat.session_input import ask_yes_no
@@ -74,31 +74,25 @@ def is_abs_path(self, attribute: attr.Attribute[Path], value: Any):
             raise ValueError(f"File_path must be an absolute path, got {value}")
 
     def _display_creation(self, prefix: str = ""):
-        ctx = SESSION_CONTEXT.get()
-
         added_lines = list[str]()
         for replacement in self.replacements:
             added_lines.extend(replacement.new_lines)
         display_information = DisplayInformation(
             self.file_path, [], added_lines, [], FileActionType.CreateFile
         )
-        ctx.stream.send(get_full_change(display_information, prefix=prefix))
+        display_full_change(display_information, prefix=prefix)
 
     def _display_deletion(self, file_lines: list[str], prefix: str = ""):
-        ctx = SESSION_CONTEXT.get()
-
         display_information = DisplayInformation(
             self.file_path,
             [],
             [],
             file_lines,
             FileActionType.DeleteFile,
         )
-        ctx.stream.send(get_full_change(display_information, prefix=prefix))
+        display_full_change(display_information, prefix=prefix)
 
     def _display_rename(self, prefix: str = ""):
-        ctx = SESSION_CONTEXT.get()
-
         display_information = DisplayInformation(
             self.file_path,
             [],
@@ -107,13 +101,11 @@ def _display_rename(self, prefix: str = ""):
             FileActionType.RenameFile,
             new_name=self.rename_file_path,
         )
-        ctx.stream.send(get_full_change(display_information, prefix=prefix))
+        display_full_change(display_information, prefix=prefix)
 
     def _display_replacement(
         self, replacement: Replacement, file_lines: list[str], prefix: str = ""
     ):
-        ctx = SESSION_CONTEXT.get()
-
         removed_block = file_lines[replacement.starting_line : replacement.ending_line]
         display_information = DisplayInformation(
             self.file_path,
@@ -125,7 +117,7 @@ def _display_replacement(
             replacement.ending_line,
             self.rename_file_path,
         )
-        ctx.stream.send(get_full_change(display_information, prefix=prefix))
+        display_full_change(display_information, prefix=prefix)
 
     def _display_replacements(self, file_lines: list[str], prefix: str = ""):
         for replacement in self.replacements: