Add support for ast.IfExp (#111)

* Add support for ast.IfExp

* Add support for recursive statement

* Unit tests

* Fix formatting issues

Author: kshxtij

src/latexify/codegen/function_codegen.py

@@ -524,6 +524,18 @@ def visit_If(self, node: ast.If) -> str:
         latex += self.visit(node)
         return latex + r", & \mathrm{otherwise} \end{array} \right."
 
+    def visit_IfExp(self, node: ast.IfExp) -> str:
+        """Visit an ifexp node"""
+        latex = r"\left\{ \begin{array}{ll} "
+        while isinstance(node, ast.IfExp):
+            cond_latex = self.visit(node.test)
+            true_latex = self.visit(node.body)
+            latex += true_latex + r", & \mathrm{if} \ " + cond_latex + r" \\ "
+            node = node.orelse
+
+        latex += self.visit(node)
+        return latex + r", & \mathrm{otherwise} \end{array} \right."
+
     def _get_sum_prod_range(self, node: ast.comprehension) -> tuple[str, str] | None:
         """Helper to process range(...) for sum and prod functions.
 

src/latexify/codegen/function_codegen_test.py

@@ -253,6 +253,38 @@ def test_visit_call_sum_prod_with_if(src_suffix: str, dest_suffix: str) -> None:
         assert FunctionCodegen().visit(node) == dest_fn + dest_suffix
 
 
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        (
+            "x if x < y else y",
+            r"\left\{ \begin{array}{ll} x,"
+            r" & \mathrm{if} \ {x < y} \\ y,"
+            r" & \mathrm{otherwise} \end{array} \right.",
+        ),
+        (
+            "x if x < y else (y if y < z else z)",
+            r"\left\{ \begin{array}{ll} x,"
+            r" & \mathrm{if} \ {x < y} \\ y,"
+            r" & \mathrm{if} \ {y < z} \\ z,"
+            r" & \mathrm{otherwise} \end{array} \right.",
+        ),
+        (
+            "x if x < y else (y if y < z else (z if z < w else w))",
+            r"\left\{ \begin{array}{ll} x,"
+            r" & \mathrm{if} \ {x < y} \\ y,"
+            r" & \mathrm{if} \ {y < z} \\ z,"
+            r" & \mathrm{if} \ {z < w} \\ w,"
+            r" & \mathrm{otherwise} \end{array} \right.",
+        ),
+    ],
+)
+def test_if_then_else(code: str, latex: str) -> None:
+    node = ast.parse(code).body[0].value
+    assert isinstance(node, ast.IfExp)
+    assert FunctionCodegen().visit(node) == latex
+
+
 @pytest.mark.parametrize(
     "code,latex",
     [