diff --git a/midas/cli/highlighter.py b/midas/cli/highlighter.py
index 1303f94..ce6c2b2 100644
--- a/midas/cli/highlighter.py
+++ b/midas/cli/highlighter.py
@@ -228,6 +228,13 @@ class PythonHighlighter(
         for item in expr.items:
             item.accept(self)
 
+    def visit_dict_expr(self, expr: p.DictExpr) -> None:
+        for key in expr.keys:
+            if key is not None:
+                key.accept(self)
+        for value in expr.values:
+            value.accept(self)
+
     def visit_subscript_expr(self, expr: p.SubscriptExpr) -> None:
         expr.object.accept(self)
         expr.index.accept(self)
@@ -240,6 +247,10 @@ class PythonHighlighter(
         if expr.step is not None:
             expr.step.accept(self)
 
+    def visit_raw_expr(self, expr: p.RawExpr) -> None: ...
+
+    def visit_raw_stmt(self, stmt: p.RawStmt) -> None: ...
+
 
 class MidasHighlighter(
     Highlighter, m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[None]
@@ -266,8 +277,9 @@ class MidasHighlighter(
     def visit_predicate_stmt(self, stmt: m.PredicateStmt) -> None:
         self.wrap(stmt, "predicate")
         self.wrap(LocatableToken(stmt.name), "predicate-name")
-        stmt.type.accept(self)
-        stmt.condition.accept(self)
+        for spec in stmt.params:
+            self._visit_param_spec(spec)
+        stmt.body.accept(self)
 
     def visit_logical_expr(self, expr: m.LogicalExpr) -> None:
         self.wrap(expr, "logical-expr")
@@ -283,6 +295,14 @@ class MidasHighlighter(
         self.wrap(expr, "unary-expr")
         expr.right.accept(self)
 
+    def visit_call_expr(self, expr: m.CallExpr) -> None:
+        self.wrap(expr, "call-expr")
+        expr.callee.accept(self)
+        for arg in expr.arguments:
+            arg.accept(self)
+        for arg in expr.keywords.values():
+            arg.accept(self)
+
     def visit_get_expr(self, expr: m.GetExpr) -> None:
         self.wrap(expr, "get-expr")
         expr.expr.accept(self)
@@ -318,8 +338,7 @@ class MidasHighlighter(
 
     def visit_function_type(self, type: m.FunctionType) -> None:
         self.wrap(type, "function")
-        for arg in type.pos_args + type.args + type.kw_args:
-            arg.type.accept(self)
+        self._visit_param_spec(type.params)
         type.returns.accept(self)
 
     def visit_extension_type(self, type: m.ExtensionType) -> None:
@@ -327,6 +346,10 @@ class MidasHighlighter(
         type.base.accept(self)
         type.extension.accept(self)
 
+    def _visit_param_spec(self, spec: m.ParamSpec) -> None:
+        for param in spec.pos + spec.mixed + spec.kw:
+            param.type.accept(self)
+
 
 class DiagnosticsHighlighter(Highlighter):
     EXTRA_CSS_PATH: Optional[Path] = Path(__file__).parent / "hl_diagnostic.css"