Supports if clause in comprehension with sum/prod. (#79)

* add some utils to analyze range.

* Support some forms in sum/prod

* support multi-clause comprehension in sum and prod

* support if in sum and prod.

* fix

Author: Yusuke Oda

src/latexify/codegen/function_codegen.py

@@ -384,23 +384,28 @@ def _get_sum_prod_info(
         scripts: list[tuple[str, str]] = []
 
         for comp in node.generators:
-            # TODO(odashi): This could be supported.
-            if comp.ifs:
-                raise exceptions.LatexifyNotSupportedError(
-                    "If-clause in comprehension is not supported."
-                )
-
             target = self.visit(comp.target)
             range_args = self._get_sum_prod_range(comp)
 
-            if range_args is not None:
+            if range_args is not None and not comp.ifs:
                 lower_rhs, upper = range_args
                 lower = f"{target} = {lower_rhs}"
             else:
                 lower_rhs = self.visit(comp.iter)
-                lower = rf"{target} \in {lower_rhs}"
+                lower_in = rf"{target} \in {lower_rhs}"
                 upper = ""
 
+                if comp.ifs:
+                    conds = [lower_in] + [self.visit(cond) for cond in comp.ifs]
+                    conds_wrapped = [r"\left(" + cond + r"\right)" for cond in conds]
+                    lower = r" \land ".join(conds_wrapped)
+                    # TODO(odashi):
+                    # Following form may be prettier, but requires amsmath.
+                    # It would be good if we have an option to switch the behavior.
+                    # lower = r"\substack{" + r" \\ ".join(lowers) + "}"
+                else:
+                    lower = lower_in
+
             scripts.append((lower, upper))
 
         return elt, scripts

src/latexify/codegen/function_codegen_test.py

@@ -108,11 +108,26 @@ def test_visit_call_sum_prod_multiple_comprehension(code: str, latex: str) -> No
     assert FunctionCodegen().visit(node) == latex
 
 
-def test_visit_call_sum_prod_with_if() -> None:
-    for fn_name in ["sum", "math.prod"]:
-        node = ast.parse(f"{fn_name}(i for y in x if y == 0)").body[0].value
-        with pytest.raises(exceptions.LatexifyNotSupportedError, match="^If-clause"):
-            FunctionCodegen().visit(node)
+@pytest.mark.parametrize(
+    "src_suffix,dest_suffix",
+    [
+        (
+            "(i for i in x if i < y)",
+            r"_{\left(i \in x\right) \land \left({i < y}\right)}^{} \left({i}\right)",
+        ),
+        (
+            "(i for i in x if i < y if f(i))",
+            r"_{\left(i \in x\right) \land \left({i < y}\right)"
+            r" \land \left(\mathrm{f}\left(i\right)\right)}^{}"
+            r" \left({i}\right)",
+        ),
+    ],
+)
+def test_visit_call_sum_prod_with_if(src_suffix: str, dest_suffix: str) -> None:
+    for src_fn, dest_fn in [("sum", r"\sum"), ("math.prod", r"\prod")]:
+        node = ast.parse(src_fn + src_suffix).body[0].value
+        assert isinstance(node, ast.Call)
+        assert FunctionCodegen().visit(node) == dest_fn + dest_suffix
 
 
 @pytest.mark.parametrize(