Improve algorithmic codegen, return multiline string

src/integration_tests/algorithmic_style_test.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import textwrap
 from typing import Any, Callable
 
 from latexify import frontend
@@ -34,17 +35,20 @@ def fact(n):
         else:
             return n * fact(n - 1)
 
-    latex = (
-        r"\begin{algorithmic}"
-        r" \Function{fact}{$n$}"
-        r" \If{$n = 0$}"
-        r" \State \Return $1$"
-        r" \Else"
-        r" \State \Return $n \mathrm{fact} \mathopen{}\left( n - 1 \mathclose{}\right)$"
-        r" \EndIf"
-        r" \EndFunction"
-        r" \end{algorithmic}"
-    )
+    latex = textwrap.dedent(
+        r"""
+        \begin{algorithmic}
+            \Function{fact}{$n$}
+                \If{$n = 0$}
+                    \State \Return $1$
+                \Else
+                    \State \Return $n \mathrm{fact} \mathopen{}\left( n - 1 \mathclose{}\right)$
+                \EndIf
+            \EndFunction
+        \end{algorithmic}
+        """  # noqa: E501
+    ).strip()
+
     check_algorithm(fact, latex)
 
 
@@ -59,19 +63,22 @@ def collatz(n):
             iterations = iterations + 1
         return iterations
 
-    latex = (
-        r"\begin{algorithmic}"
-        r" \Function{collatz}{$n$}"
-        r" \State $\mathrm{iterations} \gets 0$"
-        r" \While{$n > 1$}"
-        r" \If{$n \mathbin{\%} 2 = 0$}"
-        r" \State $n \gets \left\lfloor\frac{n}{2}\right\rfloor$"
-        r" \Else \State $n \gets 3 n + 1$"
-        r" \EndIf"
-        r" \State $\mathrm{iterations} \gets \mathrm{iterations} + 1$"
-        r" \EndWhile"
-        r" \State \Return $\mathrm{iterations}$"
-        r" \EndFunction"
-        r" \end{algorithmic}"
-    )
+    latex = textwrap.dedent(
+        r"""
+        \begin{algorithmic}
+            \Function{collatz}{$n$}
+                \State $\mathrm{iterations} \gets 0$
+                \While{$n > 1$}
+                    \If{$n \mathbin{\%} 2 = 0$}
+                        \State $n \gets \left\lfloor\frac{n}{2}\right\rfloor$
+                    \Else
+                        \State $n \gets 3 n + 1$
+                    \EndIf
+                    \State $\mathrm{iterations} \gets \mathrm{iterations} + 1$
+                \EndWhile
+                \State \Return $\mathrm{iterations}$
+            \EndFunction
+        \end{algorithmic}
+        """
+    ).strip()
     check_algorithm(collatz, latex)

src/latexify/codegen/algorithmic_codegen.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import ast
+import contextlib
+from collections.abc import Generator
 
 from latexify import exceptions
 from latexify.codegen import expression_codegen, identifier_converter
@@ -15,7 +17,10 @@ class AlgorithmicCodegen(ast.NodeVisitor):
     LaTeX expression of the given algorithm.
     """
 
+    _SPACES_PER_INDENT = 4
+
     _identifier_converter: identifier_converter.IdentifierConverter
+    _indent_level: int
 
     def __init__(
         self, *, use_math_symbols: bool = False, use_set_symbols: bool = False
@@ -33,6 +38,7 @@ def __init__(
         self._identifier_converter = identifier_converter.IdentifierConverter(
             use_math_symbols=use_math_symbols
         )
+        self._indent_level = 0
 
     def generic_visit(self, node: ast.AST) -> str:
         raise exceptions.LatexifyNotSupportedError(
@@ -46,41 +52,51 @@ def visit_Assign(self, node: ast.Assign) -> str:
         ]
         operands.append(self._expression_codegen.visit(node.value))
         operands_latex = r" \gets ".join(operands)
-        return rf"\State ${operands_latex}$"
+        return self._add_indent(rf"\State ${operands_latex}$")
 
     def visit_Expr(self, node: ast.Expr) -> str:
         """Visit an Expr node."""
-        return rf"\State ${self._expression_codegen.visit(node.value)}$"
+        return self._add_indent(
+            rf"\State ${self._expression_codegen.visit(node.value)}$"
+        )
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         """Visit a FunctionDef node."""
         # Arguments
         arg_strs = [
             self._identifier_converter.convert(arg.arg)[0] for arg in node.args.args
         ]
-        # Body
-        body_strs: list[str] = [self.visit(stmt) for stmt in node.body]
-        return (
-            rf"\begin{{algorithmic}}"
-            rf" \Function{{{node.name}}}{{${', '.join(arg_strs)}$}}"
-            f" {' '.join(body_strs)}"
-            r" \EndFunction"
-            rf" \end{{algorithmic}}"
-        )
+
+        latex = self._add_indent("\\begin{algorithmic}\n")
+        with self._increment_level():
+            latex += self._add_indent(
+                f"\\Function{{{node.name}}}{{${', '.join(arg_strs)}$}}\n"
+            )
+
+            with self._increment_level():
+                # Body
+                body_strs: list[str] = [self.visit(stmt) for stmt in node.body]
+            body_latex = "\n".join(body_strs)
+
+            latex += f"{body_latex}\n"
+            latex += self._add_indent("\\EndFunction\n")
+        return latex + self._add_indent(r"\end{algorithmic}")
 
     # TODO(ZibingZhang): support \ELSIF
     def visit_If(self, node: ast.If) -> str:
         """Visit an If node."""
         cond_latex = self._expression_codegen.visit(node.test)
-        body_latex = " ".join(self.visit(stmt) for stmt in node.body)
+        with self._increment_level():
+            body_latex = "\n".join(self.visit(stmt) for stmt in node.body)
 
-        latex = rf"\If{{${cond_latex}$}} {body_latex}"
+        latex = self._add_indent(f"\\If{{${cond_latex}$}}\n{body_latex}")
 
         if node.orelse:
-            latex += r" \Else "
-            latex += " ".join(self.visit(stmt) for stmt in node.orelse)
+            latex += "\n" + self._add_indent(r"\Else") + "\n"
+            with self._increment_level():
+                latex += "\n".join(self.visit(stmt) for stmt in node.orelse)
 
-        return latex + r" \EndIf"
+        return latex + "\n" + self._add_indent(r"\EndIf")
 
     def visit_Module(self, node: ast.Module) -> str:
         """Visit a Module node."""
@@ -89,9 +105,11 @@ def visit_Module(self, node: ast.Module) -> str:
     def visit_Return(self, node: ast.Return) -> str:
         """Visit a Return node."""
         return (
-            rf"\State \Return ${self._expression_codegen.visit(node.value)}$"
+            self._add_indent(
+                rf"\State \Return ${self._expression_codegen.visit(node.value)}$"
+            )
             if node.value is not None
-            else r"\State \Return"
+            else self._add_indent(r"\State \Return")
         )
 
     def visit_While(self, node: ast.While) -> str:
@@ -102,5 +120,25 @@ def visit_While(self, node: ast.While) -> str:
             )
 
         cond_latex = self._expression_codegen.visit(node.test)
-        body_latex = " ".join(self.visit(stmt) for stmt in node.body)
-        return rf"\While{{${cond_latex}$}} {body_latex} \EndWhile"
+        with self._increment_level():
+            body_latex = "\n".join(self.visit(stmt) for stmt in node.body)
+        return (
+            self._add_indent(f"\\While{{${cond_latex}$}}\n")
+            + f"{body_latex}\n"
+            + self._add_indent(r"\EndWhile")
+        )
+
+    @contextlib.contextmanager
+    def _increment_level(self) -> Generator[None, None, None]:
+        """Context manager controlling indent level."""
+        self._indent_level += 1
+        yield
+        self._indent_level -= 1
+
+    def _add_indent(self, line: str) -> str:
+        """Adds whitespace before the line.
+
+        Args:
+            line: The line to add whitespace to.
+        """
+        return self._indent_level * self._SPACES_PER_INDENT * " " + line

src/latexify/codegen/algorithmic_codegen_test.py

@@ -25,8 +25,14 @@ class UnknownNode(ast.AST):
 @pytest.mark.parametrize(
     "code,latex",
     [
-        ("x = 3", r"\State $x \gets 3$"),
-        ("a = b = 0", r"\State $a \gets b \gets 0$"),
+        (
+            "x = 3",
+            r"\State $x \gets 3$",
+        ),
+        (
+            "a = b = 0",
+            r"\State $a \gets b \gets 0$",
+        ),
     ],
 )
 def test_visit_assign(code: str, latex: str) -> None:
@@ -40,46 +46,65 @@ def test_visit_assign(code: str, latex: str) -> None:
     [
         (
             "def f(x): return x",
-            (
-                r"\begin{algorithmic}"
-                r" \Function{f}{$x$}"
-                r" \State \Return $x$"
-                r" \EndFunction"
-                r" \end{algorithmic}"
-            ),
+            r"""
+            \begin{algorithmic}
+                \Function{f}{$x$}
+                    \State \Return $x$
+                \EndFunction
+            \end{algorithmic}
+            """,
         ),
         (
             "def xyz(a, b, c): return 3",
-            (
-                r"\begin{algorithmic}"
-                r" \Function{xyz}{$a, b, c$}"
-                r" \State \Return $3$"
-                r" \EndFunction"
-                r" \end{algorithmic}"
-            ),
+            r"""
+            \begin{algorithmic}
+                \Function{xyz}{$a, b, c$}
+                    \State \Return $3$
+                \EndFunction
+            \end{algorithmic}
+            """,
         ),
     ],
 )
 def test_visit_functiondef(code: str, latex: str) -> None:
     node = ast.parse(textwrap.dedent(code)).body[0]
     assert isinstance(node, ast.FunctionDef)
-    assert algorithmic_codegen.AlgorithmicCodegen().visit(node) == latex
+    assert (
+        algorithmic_codegen.AlgorithmicCodegen().visit(node)
+        == textwrap.dedent(latex).strip()
+    )
 
 
 @pytest.mark.parametrize(
     "code,latex",
     [
-        ("if x < y: return x", r"\If{$x < y$} \State \Return $x$ \EndIf"),
+        (
+            "if x < y: return x",
+            r"""
+            \If{$x < y$}
+                \State \Return $x$
+            \EndIf
+            """,
+        ),
         (
             "if True: x\nelse: y",
-            r"\If{$\mathrm{True}$} \State $x$ \Else \State $y$ \EndIf",
+            r"""
+            \If{$\mathrm{True}$}
+                \State $x$
+            \Else
+                \State $y$
+            \EndIf
+            """,
         ),
     ],
 )
 def test_visit_if(code: str, latex: str) -> None:
-    node = ast.parse(textwrap.dedent(code)).body[0]
+    node = ast.parse(code).body[0]
     assert isinstance(node, ast.If)
-    assert algorithmic_codegen.AlgorithmicCodegen().visit(node) == latex
+    assert (
+        algorithmic_codegen.AlgorithmicCodegen().visit(node)
+        == textwrap.dedent(latex).strip()
+    )
 
 
 @pytest.mark.parametrize(
@@ -96,7 +121,7 @@ def test_visit_if(code: str, latex: str) -> None:
     ],
 )
 def test_visit_return(code: str, latex: str) -> None:
-    node = ast.parse(textwrap.dedent(code)).body[0]
+    node = ast.parse(code).body[0]
     assert isinstance(node, ast.Return)
     assert algorithmic_codegen.AlgorithmicCodegen().visit(node) == latex
 
@@ -106,14 +131,21 @@ def test_visit_return(code: str, latex: str) -> None:
     [
         (
             "while x < y: x = x + 1",
-            r"\While{$x < y$} \State $x \gets x + 1$ \EndWhile",
+            r"""
+            \While{$x < y$}
+                \State $x \gets x + 1$
+            \EndWhile
+            """,
         )
     ],
 )
 def test_visit_while(code: str, latex: str) -> None:
-    node = ast.parse(textwrap.dedent(code)).body[0]
+    node = ast.parse(code).body[0]
     assert isinstance(node, ast.While)
-    assert algorithmic_codegen.AlgorithmicCodegen().visit(node) == latex
+    assert (
+        algorithmic_codegen.AlgorithmicCodegen().visit(node)
+        == textwrap.dedent(latex).strip()
+    )
 
 
 def test_visit_while_with_else() -> None: