From 25bd895dde07cfa99c5cf93ac58438ea1739cc59 Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Sun, 7 Jun 2026 13:42:15 +0200
Subject: [PATCH] feat(cli): improve diagnostic printing

---
 midas/checker/diagnostic.py |  8 +++--
 midas/cli/ansi.py           | 41 ++++++++++++++++++++++++++
 midas/cli/main.py           | 58 +++++++++++++++++++++++++++++++++++--
 3 files changed, 103 insertions(+), 4 deletions(-)
 create mode 100644 midas/cli/ansi.py

diff --git a/midas/checker/diagnostic.py b/midas/checker/diagnostic.py
index 45514b4..77f687e 100644
--- a/midas/checker/diagnostic.py
+++ b/midas/checker/diagnostic.py
@@ -19,7 +19,8 @@ class Diagnostic:
     type: DiagnosticType
     message: str
 
-    def __str__(self) -> str:
+    @property
+    def location_str(self) -> str:
         start_loc: str = f"L{self.location.lineno}:{self.location.col_offset+1}"
         end_loc: Optional[str] = ""
         if (
@@ -30,4 +31,7 @@ class Diagnostic:
         loc: str = (
             f"at {start_loc}" if end_loc is None else f"from {start_loc} to {end_loc}"
         )
-        return f"{self.type} in {self.file_path} {loc}: {self.message}"
+        return f"{self.type} in {self.file_path} {loc}"
+
+    def __str__(self) -> str:
+        return f"{self.location_str}: {self.message}"
diff --git a/midas/cli/ansi.py b/midas/cli/ansi.py
new file mode 100644
index 0000000..f929be4
--- /dev/null
+++ b/midas/cli/ansi.py
@@ -0,0 +1,41 @@
+class Ansi:
+    CTRL = "\x1b["
+    RESET = CTRL + "0m"
+    BOLD = CTRL + "1m"
+    DIM = CTRL + "2m"
+    ITALIC = CTRL + "3m"
+    UNDERLINE = CTRL + "4m"
+
+    BLACK = 0
+    RED = 1
+    GREEN = 2
+    YELLOW = 3
+    BLUE = 4
+    MAGENTA = 5
+    CYAN = 6
+    WHITE = 7
+
+    BRIGHT_BLACK = 60
+    BRIGHT_RED = 61
+    BRIGHT_GREEN = 62
+    BRIGHT_YELLOW = 63
+    BRIGHT_BLUE = 64
+    BRIGHT_MAGENTA = 65
+    BRIGHT_CYAN = 66
+    BRIGHT_WHITE = 67
+
+    @classmethod
+    def FG(cls, col: int) -> str:
+        return f"{cls.CTRL}{30 + col}m"
+
+    @classmethod
+    def BG(cls, col: int) -> str:
+        return f"{cls.CTRL}{40 + col}m"
+
+    @classmethod
+    def FG_RGB(cls, r: int, g: int, b: int) -> str:
+        return f"{cls.CTRL}38;2;{r};{g};{b}m"
+
+    @classmethod
+    def BG_RGB(cls, r: int, g: int, b: int) -> str:
+        return f"{cls.CTRL}48;2;{r};{g};{b}m"
diff --git a/midas/cli/main.py b/midas/cli/main.py
index 95b6ca1..ae4295b 100644
--- a/midas/cli/main.py
+++ b/midas/cli/main.py
@@ -8,10 +8,12 @@ import click
 
 import midas.ast.midas as m
 import midas.ast.python as p
+from midas.ast.location import Location
 from midas.ast.printer import MidasAstPrinter, MidasPrinter, PythonAstPrinter
 from midas.checker.checker import Checker
-from midas.checker.diagnostic import Diagnostic
+from midas.checker.diagnostic import Diagnostic, DiagnosticType
 from midas.checker.types import Type
+from midas.cli.ansi import Ansi
 from midas.cli.highlighter import (
     DiagnosticsHighlighter,
     Highlighter,
@@ -32,6 +34,57 @@ def midas():
     pass
 
 
+def print_diagnostic(lines: list[str], diagnostic: Diagnostic, indent: int = 4):
+    """Pretty-print a diagnostic, showing some context if possible
+
+    If the diagnostic concerns a specific part of one line, the line is shown
+    with the affected part highlighted. The message is clearly printed under the
+    line with an underline further indicating the target expression.
+
+    If multiple lines are concerned, no context is shown, only the
+    diagnostic type, location and message
+
+    Args:
+        lines (list[str]): source code lines
+        diagnostic (Diagnostic): the diagnostic to print
+        indent (int, optional): the number of spaces added before the target line to indent if from the location header. Defaults to 4.
+    """
+
+    loc: Location = diagnostic.location
+    if loc.lineno != loc.end_lineno:
+        print(diagnostic)
+        return
+
+    start_offset: int = loc.col_offset
+    end_offset: int = loc.end_col_offset or (start_offset + 1)
+
+    line: str = lines[loc.lineno - 1]
+    before: str = line[:start_offset]
+    after: str = line[end_offset:]
+
+    color: int = {
+        DiagnosticType.ERROR: Ansi.RED,
+        DiagnosticType.WARNING: Ansi.YELLOW,
+        DiagnosticType.INFO: Ansi.CYAN,
+    }.get(diagnostic.type, Ansi.WHITE)
+
+    subject: str = Ansi.FG(color) + line[start_offset:end_offset] + Ansi.RESET
+    cursor: str = (
+        " " * start_offset
+        + Ansi.FG(color)
+        + "~" * (end_offset - start_offset)
+        + "> "
+        + diagnostic.message
+        + Ansi.RESET
+    )
+
+    indent_str: str = " " * indent
+    print(diagnostic.location_str + ":")
+    print(indent_str + before + subject + after)
+    print(indent_str + cursor)
+    print()
+
+
 @midas.command()
 @click.option("-l", "--highlight", type=click.File("w"))
 @click.option("-t", "--types", type=click.File("r"), multiple=True)
@@ -57,8 +110,9 @@ def compile(
         types_paths=types_paths,
     )
     diagnostics: list[Diagnostic] = checker.check(stmts)
+    lines: list[str] = source.split("\n")
     for diagnostic in diagnostics:
-        print(diagnostic)
+        print_diagnostic(lines, diagnostic)
 
     if verbose:
         print(