MatchOr and guard statements

src/latexify/codegen/function_codegen.py

@@ -140,7 +140,7 @@ def visit_If(self, node: ast.If) -> str:
         return latex + r", & \mathrm{otherwise} \end{array} \right."
 
     def visit_Match(self, node: ast.Match) -> str:
-        """Visit a Match node"""
+        """Visit a Match node."""
         if not (
             len(node.cases) >= 2
             and isinstance(node.cases[-1].pattern, ast.MatchAs)
@@ -162,21 +162,47 @@ def visit_Match(self, node: ast.Match) -> str:
             if i < len(node.cases) - 1:
                 body_latex = self.visit(case.body[0])
                 cond_latex = self.visit(case.pattern)
-                case_latexes.append(
-                    body_latex + r", & \mathrm{if} \ " + subject_latex + cond_latex
-                )
+                if case.guard is not None:
+                    cond_latex = self._expression_codegen.visit(case.guard)
+
+                case_latexes.append(body_latex + r", & \mathrm{if} \ " + cond_latex)
             else:
                 case_latexes.append(
-                    self.visit(node.cases[-1].body[0]) + r", & \mathrm{otherwise}"
+                    self.visit(case.body[0]) + r", & \mathrm{otherwise}"
                 )
 
-        return (
-            r"\left\{ \begin{array}{ll} "
+        latex = (
+            r"\left\{ \begin{array}{ll}"
             + r" \\ ".join(case_latexes)
             + r" \end{array} \right."
         )
 
+        latex_final = latex.replace("subject_name", subject_latex)
+        return latex_final
+
     def visit_MatchValue(self, node: ast.MatchValue) -> str:
-        """Visit a MatchValue node"""
+        """Visit a MatchValue node."""
         latex = self._expression_codegen.visit(node.value)
-        return " = " + latex
+        return "subject_name = " + latex
+
+    def visit_MatchAs(self, node: ast.MatchAs) -> str:
+        """
+        Visit a MatchAs node.
+        If MatchAs is a wildcard, return 'otherwise' case, otherwise throw an error.
+        """
+        if node.pattern is None:
+            return ""
+        else:
+            raise exceptions.LatexifySyntaxError(
+                "Nonempty as-patterns are not supported in MatchAs nodes."
+            )
+
+    def visit_MatchOr(self, node: ast.MatchOr) -> str:
+        """Visit a MatchOr node."""
+        case_latexes = []
+        for i, pattern in enumerate(node.patterns):
+            if i == 0:
+                case_latexes.append(self.visit(pattern))
+            else:
+                case_latexes.append(r" \lor " + self.visit(pattern))
+        return "".join(case_latexes)

src/latexify/codegen/function_codegen_match_test.py

@@ -12,7 +12,7 @@
 
 
 @test_utils.require_at_least(10)
-def test_functiondef_match() -> None:
+def test_visit_functiondef_match() -> None:
     tree = ast.parse(
         textwrap.dedent(
             """
@@ -28,15 +28,15 @@ def f(x):
     expected = (
         r"f(x) ="
         r" \left\{ \begin{array}{ll}"
-        r" 1, & \mathrm{if} \ x = 0 \\"
+        r"1, & \mathrm{if} \ x = 0 \\"
         r" 3 x, & \mathrm{otherwise}"
         r" \end{array} \right."
     )
     assert function_codegen.FunctionCodegen().visit(tree) == expected
 
 
 @test_utils.require_at_least(10)
-def test_matchvalue() -> None:
+def test_visit_match() -> None:
     tree = ast.parse(
         textwrap.dedent(
             """
@@ -50,15 +50,15 @@ def test_matchvalue() -> None:
     ).body[0]
     expected = (
         r"\left\{ \begin{array}{ll}"
-        r" 1, & \mathrm{if} \ x = 0 \\"
+        r"1, & \mathrm{if} \ x = 0 \\"
         r" 2, & \mathrm{otherwise}"
         r" \end{array} \right."
     )
     assert function_codegen.FunctionCodegen().visit(tree) == expected
 
 
 @test_utils.require_at_least(10)
-def test_multiple_matchvalue() -> None:
+def test_visit_multiple_match_cases() -> None:
     tree = ast.parse(
         textwrap.dedent(
             """
@@ -74,7 +74,7 @@ def test_multiple_matchvalue() -> None:
     ).body[0]
     expected = (
         r"\left\{ \begin{array}{ll}"
-        r" 1, & \mathrm{if} \ x = 0 \\"
+        r"1, & \mathrm{if} \ x = 0 \\"
         r" 2, & \mathrm{if} \ x = 1 \\"
         r" 3, & \mathrm{otherwise}"
         r" \end{array} \right."
@@ -83,7 +83,7 @@ def test_multiple_matchvalue() -> None:
 
 
 @test_utils.require_at_least(10)
-def test_single_matchvalue_no_wildcards() -> None:
+def test_visit_single_match_case_no_wildcards() -> None:
     tree = ast.parse(
         textwrap.dedent(
             """
@@ -102,7 +102,7 @@ def test_single_matchvalue_no_wildcards() -> None:
 
 
 @test_utils.require_at_least(10)
-def test_multiple_matchvalue_no_wildcards() -> None:
+def test_visit_multiple_match_cases_no_wildcards() -> None:
     tree = ast.parse(
         textwrap.dedent(
             """
@@ -123,7 +123,7 @@ def test_multiple_matchvalue_no_wildcards() -> None:
 
 
 @test_utils.require_at_least(10)
-def test_matchas_nonempty() -> None:
+def test_visit_only_wildcard_in_matchas() -> None:
     tree = ast.parse(
         textwrap.dedent(
             """
@@ -137,14 +137,14 @@ def test_matchas_nonempty() -> None:
     ).body[0]
 
     with pytest.raises(
-        exceptions.LatexifyNotSupportedError,
-        match=r"^Unsupported AST: MatchAs$",
+        exceptions.LatexifySyntaxError,
+        match=r"^Nonempty as-patterns are not supported in MatchAs nodes.$",
     ):
         function_codegen.FunctionCodegen().visit(tree)
 
 
 @test_utils.require_at_least(10)
-def test_matchvalue_no_return() -> None:
+def test_visit_match_case_no_return() -> None:
     tree = ast.parse(
         textwrap.dedent(
             """
@@ -165,7 +165,7 @@ def test_matchvalue_no_return() -> None:
 
 
 @test_utils.require_at_least(10)
-def test_matchvalue_mutliple_statements() -> None:
+def test_visit_match_case_mutliple_statements() -> None:
     tree = ast.parse(
         textwrap.dedent(
             """
@@ -184,3 +184,112 @@ def test_matchvalue_mutliple_statements() -> None:
         match=r"^Match cases must contain exactly 1 return statement\.$",
     ):
         function_codegen.FunctionCodegen().visit(tree)
+
+
+@test_utils.require_at_least(10)
+def test_visit_match_case_with_if() -> None:
+    tree = ast.parse(
+        textwrap.dedent(
+            """
+            match x:
+                case x if x > 0:
+                    return 1
+                case _:
+                    return 2
+            """
+        )
+    ).body[0]
+
+    assert (
+        function_codegen.FunctionCodegen().visit(tree)
+        == r"\left\{ \begin{array}{ll}"
+        + r"1, & \mathrm{if} \ x > 0 \\ "
+        + r"2, & \mathrm{otherwise} \end{array} \right."
+    )
+
+
+@test_utils.require_at_least(10)
+def test_visit_match_case_with_if_and() -> None:
+    tree = ast.parse(
+        textwrap.dedent(
+            """
+            match x:
+                case x if x > 0 and x <= 10:
+                    return 1
+                case _:
+                    return 2
+            """
+        )
+    ).body[0]
+
+    assert (
+        function_codegen.FunctionCodegen().visit(tree)
+        == r"\left\{ \begin{array}{ll}1, & \mathrm{if} "
+        + r"\ x > 0 \land x \le 10 \\"
+        + r" 2, & \mathrm{otherwise} \end{array} \right."
+    )
+
+
+@test_utils.require_at_least(10)
+def test_visit_matchcase_with_if_or() -> None:
+    tree = ast.parse(
+        textwrap.dedent(
+            """
+            match x:
+                case x if x > 0 or x <= 10:
+                    return 1
+                case _:
+                    return 2
+            """
+        )
+    ).body[0]
+
+    assert (
+        function_codegen.FunctionCodegen().visit(tree)
+        == r"\left\{ \begin{array}{ll}1,"
+        + r" & \mathrm{if} \ x > 0 \lor x \le 10 \\"
+        + r" 2, & \mathrm{otherwise} \end{array} \right."
+    )
+
+
+@test_utils.require_at_least(10)
+def test_visit_match_case_with_combined_condition() -> None:
+    tree = ast.parse(
+        textwrap.dedent(
+            """
+            match x:
+                case x if 0 < x <= 10:
+                    return 1
+                case _:
+                    return 2
+            """
+        )
+    ).body[0]
+
+    assert (
+        function_codegen.FunctionCodegen().visit(tree)
+        == r"\left\{ \begin{array}{ll}1,"
+        + r" & \mathrm{if} \ 0 < x \le 10 \\ 2,"
+        + r" & \mathrm{otherwise} \end{array} \right."
+    )
+
+
+@test_utils.require_at_least(10)
+def test_visit_match_case_or() -> None:
+    tree = ast.parse(
+        textwrap.dedent(
+            """
+            match x:
+                case 0 | 1:
+                    return 1
+                case _:
+                    return 2
+            """
+        )
+    ).body[0]
+
+    assert (
+        function_codegen.FunctionCodegen().visit(tree)
+        == r"\left\{ \begin{array}{ll}1, & \mathrm{if} \ x = 0 \lor x = 1 \\"
+        + r" 2, & \mathrm{otherwise} \end{array} \right."
+    )