From bea3f399adfb37a19fde7760ef00883e58242de3 Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Mon, 1 Jun 2026 15:02:12 +0200
Subject: [PATCH] feat(checker): handle ternary expression

---
 .../01_simple_type_checking/03_control_flow.py |  4 +++-
 midas/checker/checker.py                       | 18 +++++++++++++++++-
 midas/cli/highlighter.py                       |  2 ++
 midas/resolver/resolver.py                     |  5 +++++
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/examples/01_simple_type_checking/03_control_flow.py b/examples/01_simple_type_checking/03_control_flow.py
index 07e11a8..90f5530 100644
--- a/examples/01_simple_type_checking/03_control_flow.py
+++ b/examples/01_simple_type_checking/03_control_flow.py
@@ -11,4 +11,6 @@ c = minimum(a, b)
 def factorial(n: int) -> int:
     if n <= 1:
         return 1
-    return n * factorial(n - 1)
\ No newline at end of file
+    return n * factorial(n - 1)
+
+category = "Category 1" if a < 10 else "Category 2"
\ No newline at end of file
diff --git a/midas/checker/checker.py b/midas/checker/checker.py
index 603f0be..c873863 100644
--- a/midas/checker/checker.py
+++ b/midas/checker/checker.py
@@ -9,7 +9,7 @@ from midas.ast.location import Location
 from midas.checker.diagnostic import Diagnostic, DiagnosticType
 from midas.checker.environment import Environment
 from midas.checker.operators import COMPARATOR_METHODS, OPERATOR_METHODS
-from midas.checker.types import Function, Type, UnitType, UnknownType
+from midas.checker.types import BaseType, Function, SimpleType, Type, UnitType, UnknownType
 from midas.lexer.midas import MidasLexer
 from midas.lexer.token import Token
 from midas.parser.midas import MidasParser
@@ -405,6 +405,22 @@ class Checker(
     def visit_cast_expr(self, expr: p.CastExpr) -> Type:
         return expr.type.accept(self)
 
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> Type:
+        test_type: Type = expr.test.accept(self)
+
+        # TODO Allow subtypes or any type
+        if test_type != self.ctx.get_type("bool"):
+            self.error(
+                expr.test.location, f"If test must be a boolean, got {test_type}"
+            )
+
+        true_type: Type = expr.if_true.accept(self)
+        false_type: Type = expr.if_false.accept(self)
+        if true_type != false_type:
+            self.error(expr.location, f"Type mismatch in ternary if branches: true={true_type} != false={false_type}")
+            return UnknownType()
+        return true_type
+
     def visit_base_type(self, node: p.BaseType) -> Type:
         return self.ctx.get_type(node.base)
 
diff --git a/midas/cli/highlighter.py b/midas/cli/highlighter.py
index bce4129..c0ecea7 100644
--- a/midas/cli/highlighter.py
+++ b/midas/cli/highlighter.py
@@ -203,6 +203,8 @@ class PythonHighlighter(
 
     def visit_cast_expr(self, expr: p.CastExpr) -> None: ...
 
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> None: ...
+
 
 class MidasHighlighter(Highlighter, m.Stmt.Visitor[None], m.Expr.Visitor[None]):
     EXTRA_CSS_PATH: Optional[Path] = Path(__file__).parent / "hl_midas.css"
diff --git a/midas/resolver/resolver.py b/midas/resolver/resolver.py
index 221c221..15166bd 100644
--- a/midas/resolver/resolver.py
+++ b/midas/resolver/resolver.py
@@ -180,3 +180,8 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
 
     def visit_cast_expr(self, expr: p.CastExpr) -> None:
         self.resolve(expr.expr)
+
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> None:
+        self.resolve(expr.test)
+        self.resolve(expr.if_true)
+        self.resolve(expr.if_false)