From bbd0e3ae8d73cd034f6714af0639e29225cc8e5f Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Mon, 25 May 2026 22:14:44 +0200
Subject: [PATCH] feat(cli): update highlighter with new nodes

---
 midas/cli/highlighter.py | 47 ++++++++++++++++++++++++++++++++--------
 midas/cli/main.py        | 24 ++++++--------------
 2 files changed, 45 insertions(+), 26 deletions(-)

diff --git a/midas/cli/highlighter.py b/midas/cli/highlighter.py
index 051ae79..45ed55c 100644
--- a/midas/cli/highlighter.py
+++ b/midas/cli/highlighter.py
@@ -104,7 +104,12 @@ class Highlighter(ABC):
                 self.openings.setdefault((l + 1, 0), []).append(opening)
 
 
-class PythonHighlighter(Highlighter, p.Expr.Visitor[None]):
+class PythonHighlighter(
+    Highlighter,
+    p.MidasType.Visitor[None],
+    p.Stmt.Visitor[None],
+    p.Expr.Visitor[None],
+):
     EXTRA_CSS_PATH: Optional[Path] = Path(__file__).parent / "hl_python.css"
 
     def highlight(self, node: Highlightable[PythonHighlighter]):
@@ -130,15 +135,39 @@ class PythonHighlighter(Highlighter, p.Expr.Visitor[None]):
         for column in node.columns:
             column.accept(self)
 
-    def visit_function(self, node: p.Function) -> None:
-        self.wrap(node, "function")
-        for arg in node.posonlyargs + node.args + node.kwonlyargs:
-            arg.accept(self)
+    def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
+        stmt.expr.accept(self)
 
-    def visit_function_argument(self, node: p.FunctionArgument) -> None:
-        self.wrap(node, "argument")
-        if node.type is not None:
-            node.type.accept(self)
+    def visit_function(self, stmt: p.Function) -> None:
+        self.wrap(stmt, "function")
+        for arg in stmt.posonlyargs + stmt.args + stmt.kwonlyargs:
+            self._highlight_function_argument(arg)
+
+    def _highlight_function_argument(self, arg: p.Function.Argument) -> None:
+        self.wrap(arg, "argument")
+        if arg.type is not None:
+            arg.type.accept(self)
+
+    def visit_type_assign(self, stmt: p.TypeAssign) -> None:
+        stmt.type.accept(self)
+
+    def visit_assign_expr(self, expr: p.AssignExpr) -> None: ...
+
+    def visit_binary_expr(self, expr: p.BinaryExpr) -> None: ...
+
+    def visit_unary_expr(self, expr: p.UnaryExpr) -> None: ...
+
+    def visit_call_expr(self, expr: p.CallExpr) -> None: ...
+
+    def visit_get_expr(self, expr: p.GetExpr) -> None: ...
+
+    def visit_literal_expr(self, expr: p.LiteralExpr) -> None: ...
+
+    def visit_variable_expr(self, expr: p.VariableExpr) -> None: ...
+
+    def visit_logical_expr(self, expr: p.LogicalExpr) -> None: ...
+
+    def visit_set_expr(self, expr: p.SetExpr) -> None: ...
 
 
 class MidasHighlighter(Highlighter, m.Stmt.Visitor[None], m.Expr.Visitor[None]):
diff --git a/midas/cli/main.py b/midas/cli/main.py
index 7319226..fb0e716 100644
--- a/midas/cli/main.py
+++ b/midas/cli/main.py
@@ -5,6 +5,7 @@ import click
 
 from midas.ast.location import Location
 import midas.ast.midas as m
+import midas.ast.python as p
 from midas.ast.printer import PythonAstPrinter
 from midas.cli.highlighter import Highlighter, MidasHighlighter, PythonHighlighter
 from midas.lexer.midas import MidasLexer
@@ -40,21 +41,13 @@ def dump_ast(output: Optional[TextIO], parse: bool, file: TextIO):
 
     if parse:
         parser = PythonParser()
-        parser.visit(tree)
+        stmts: list[p.Stmt] = parser.parse_module(tree)
         printer = PythonAstPrinter()
         dump = ""
-        for name, annotation in parser.annotations:
-            dump += f"{name} = "
-            if annotation is None:
-                dump += "None"
-            else:
-                dump += printer.print(annotation)
+        for stmt in stmts:
+            dump += printer.print(stmt)
             dump += "\n"
 
-        dump += "\n# Functions\n\n"
-
-        for func in parser.functions:
-            dump += printer.print(func) + "\n"
     else:
         dump = ast.dump(tree, indent=4)
 
@@ -67,13 +60,10 @@ def dump_ast(output: Optional[TextIO], parse: bool, file: TextIO):
 def highlight_python(source: str, path: str) -> Highlighter:
     tree: ast.Module = ast.parse(source, filename=path)
     parser = PythonParser()
-    parser.visit(tree)
+    stmts: list[p.Stmt] = parser.parse_module(tree)
     highlighter = PythonHighlighter(source)
-    for _, annotation in parser.annotations:
-        if annotation is not None:
-            highlighter.highlight(annotation)
-    for func in parser.functions:
-        highlighter.highlight(func)
+    for stmt in stmts:
+        highlighter.highlight(stmt)
     return highlighter