Improve visit_Compare implementation (#61)

* improve visit_Compare

* support multiple comparators

Author: Yusuke Oda

src/integration_tests/regression_test.py

@@ -38,7 +38,7 @@ def sinc(x):
 
 sinc_latex = (
     r"\mathrm{sinc}(x) \triangleq \left\{ \begin{array}{ll} 1, & \mathrm{if} \ "
-    r"x=0 \\ \frac{\sin{\left({x}\right)}}{x}, & \mathrm{otherwise} \end{array}"
+    r"{x = 0} \\ \frac{\sin{\left({x}\right)}}{x}, & \mathrm{otherwise} \end{array}"
     r" \right."
 )
 

src/latexify/latexify_visitor.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import ast
+from typing import ClassVar
 
 from latexify import constants
 from latexify import math_symbols
@@ -246,27 +247,26 @@ def _wrap(child):
             return reprs[type(node.op)]()
         return r"\mathrm{unknown\_binop}(" + _unwrap(lhs) + ", " + _unwrap(rhs) + ")"
 
-    def visit_Compare(self, node, action):  # pylint: disable=invalid-name
+    _compare_ops: ClassVar[dict[type[ast.cmpop], str]] = {
+        ast.Eq: "=",
+        ast.Gt: ">",
+        ast.GtE: r"\ge",
+        ast.In: r"\in",
+        ast.Is: r"\equiv",
+        ast.IsNot: r"\not\equiv",
+        ast.Lt: "<",
+        ast.LtE: r"\le",
+        ast.NotEq: r"\ne",
+        ast.NotIn: r"\notin",
+    }
+
+    def visit_Compare(self, node: ast.Compare, action):  # pylint: disable=invalid-name
         """Visit a compare node."""
-        lstr = self.visit(node.left)
-        rstr = self.visit(node.comparators[0])
-
-        if isinstance(node.ops[0], ast.Eq):
-            return lstr + "=" + rstr
-        if isinstance(node.ops[0], ast.Gt):
-            return lstr + ">" + rstr
-        if isinstance(node.ops[0], ast.Lt):
-            return lstr + "<" + rstr
-        if isinstance(node.ops[0], ast.GtE):
-            return lstr + r"\ge " + rstr
-        if isinstance(node.ops[0], ast.LtE):
-            return lstr + r"\le " + rstr
-        if isinstance(node.ops[0], ast.NotEq):
-            return lstr + r"\ne " + rstr
-        if isinstance(node.ops[0], ast.Is):
-            return lstr + r"\equiv" + rstr
-
-        return r"\mathrm{unknown\_comparator}(" + lstr + ", " + rstr + ")"
+        lhs = self.visit(node.left)
+        ops = [self._compare_ops[type(x)] for x in node.ops]
+        rhs = [self.visit(x) for x in node.comparators]
+        ops_rhs = [f" {o} {r}" for o, r in zip(ops, rhs)]
+        return "{" + lhs + "".join(ops_rhs) + "}"
 
     def visit_BoolOp(self, node, action):  # pylint: disable=invalid-name
         logic_operator = (

src/latexify/latexify_visitor_test.py

@@ -0,0 +1,46 @@
+"""Tests for latexify.latexify_visitor."""
+
+import ast
+import pytest
+
+from latexify.latexify_visitor import LatexifyVisitor
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        # 1 comparator
+        ("a == b", "{a = b}"),
+        ("a > b", "{a > b}"),
+        ("a >= b", r"{a \ge b}"),
+        ("a in b", r"{a \in b}"),
+        ("a is b", r"{a \equiv b}"),
+        ("a is not b", r"{a \not\equiv b}"),
+        ("a < b", "{a < b}"),
+        ("a <= b", r"{a \le b}"),
+        ("a != b", r"{a \ne b}"),
+        ("a not in b", r"{a \notin b}"),
+        # 2 comparators
+        ("a == b == c", "{a = b = c}"),
+        ("a == b > c", "{a = b > c}"),
+        ("a == b >= c", r"{a = b \ge c}"),
+        ("a == b < c", "{a = b < c}"),
+        ("a == b <= c", r"{a = b \le c}"),
+        ("a > b == c", "{a > b = c}"),
+        ("a > b > c", "{a > b > c}"),
+        ("a > b >= c", r"{a > b \ge c}"),
+        ("a >= b == c", r"{a \ge b = c}"),
+        ("a >= b > c", r"{a \ge b > c}"),
+        ("a >= b >= c", r"{a \ge b \ge c}"),
+        ("a < b == c", "{a < b = c}"),
+        ("a < b < c", "{a < b < c}"),
+        ("a < b <= c", r"{a < b \le c}"),
+        ("a <= b == c", r"{a \le b = c}"),
+        ("a <= b < c", r"{a \le b < c}"),
+        ("a <= b <= c", r"{a \le b \le c}"),
+    ],
+)
+def test_visit_compare(code: str, latex: str) -> None:
+    tree = ast.parse(code).body[0].value
+    assert isinstance(tree, ast.Compare)
+    assert LatexifyVisitor().visit(tree) == latex