From a15e7f4d1e75cffd92e1307080bcd618e7fbfc96 Mon Sep 17 00:00:00 2001 From: Zibing Zhang Date: Mon, 12 Dec 2022 23:29:06 +0000 Subject: [PATCH 1/6] better string output --- .../algorithmic_style_test.py | 60 ++++++++------ src/latexify/codegen/algorithmic_codegen.py | 65 +++++++++++---- .../codegen/algorithmic_codegen_test.py | 82 +++++++++++++------ 3 files changed, 138 insertions(+), 69 deletions(-) diff --git a/src/integration_tests/algorithmic_style_test.py b/src/integration_tests/algorithmic_style_test.py index 047d6f9..8ffdeaf 100644 --- a/src/integration_tests/algorithmic_style_test.py +++ b/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,21 @@ 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}" - ) + # noqa: E501 + 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 +64,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) diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py index 2fcb16b..de951e3 100644 --- a/src/latexify/codegen/algorithmic_codegen.py +++ b/src/latexify/codegen/algorithmic_codegen.py @@ -15,7 +15,10 @@ class AlgorithmicCodegen(ast.NodeVisitor): LaTeX expression of the given algorithm. """ + _SPACES_PER_INDENT = 4 + _identifier_converter: identifier_converter.IdentifierConverter + _indent: int def __init__( self, *, use_math_symbols: bool = False, use_set_symbols: bool = False @@ -33,6 +36,7 @@ def __init__( self._identifier_converter = identifier_converter.IdentifierConverter( use_math_symbols=use_math_symbols ) + self._indent = 0 def generic_visit(self, node: ast.AST) -> str: raise exceptions.LatexifyNotSupportedError( @@ -46,11 +50,11 @@ 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 rf"{self._prefix()}\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 rf"{self._prefix()}\State ${self._expression_codegen.visit(node.value)}$" def visit_FunctionDef(self, node: ast.FunctionDef) -> str: """Visit a FunctionDef node.""" @@ -58,29 +62,42 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str: arg_strs = [ self._identifier_converter.convert(arg.arg)[0] for arg in node.args.args ] + + latex = f"{self._prefix()}\\begin{{algorithmic}}\n" + + self._indent += 1 + latex += ( + f"{self._prefix()}\\Function{{{node.name}}}{{${', '.join(arg_strs)}$}}\n" + ) + # Body + self._indent += 1 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}}" - ) + self._indent -= 1 + body_latex = "\n".join(body_strs) + + latex += f"{body_latex}\n{self._prefix()}\\EndFunction\n" + self._indent -= 1 + + return latex + rf"{self._prefix()}\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) + self._indent += 1 + body_latex = "\n".join(self.visit(stmt) for stmt in node.body) + self._indent -= 1 - latex = rf"\If{{${cond_latex}$}} {body_latex}" + latex = f"{self._prefix()}\\If{{${cond_latex}$}}\n{body_latex}" if node.orelse: - latex += r" \Else " - latex += " ".join(self.visit(stmt) for stmt in node.orelse) + latex += f"\n{self._prefix()}\\Else\n" + self._indent += 1 + latex += "\n".join(self.visit(stmt) for stmt in node.orelse) + self._indent -= 1 - return latex + r" \EndIf" + return latex + f"\n{self._prefix()}\\EndIf" def visit_Module(self, node: ast.Module) -> str: """Visit a Module node.""" @@ -89,9 +106,12 @@ 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)}$" + ( + rf"{self._prefix()}\State \Return" + f" ${self._expression_codegen.visit(node.value)}$" + ) if node.value is not None - else r"\State \Return" + else rf"{self._prefix()}\State \Return" ) def visit_While(self, node: ast.While) -> str: @@ -102,5 +122,14 @@ 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" + self._indent += 1 + body_latex = "\n".join(self.visit(stmt) for stmt in node.body) + self._indent -= 1 + return ( + f"{self._prefix()}\\While{{${cond_latex}$}}\n" + f"{body_latex}\n" + rf"{self._prefix()}\EndWhile" + ) + + def _prefix(self) -> str: + return self._indent * self._SPACES_PER_INDENT * " " diff --git a/src/latexify/codegen/algorithmic_codegen_test.py b/src/latexify/codegen/algorithmic_codegen_test.py index a972beb..80a2d0c 100644 --- a/src/latexify/codegen/algorithmic_codegen_test.py +++ b/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: From 72457f768573809f7aea08bfcc16e1df754437ec Mon Sep 17 00:00:00 2001 From: Zibing Zhang Date: Mon, 12 Dec 2022 23:34:15 +0000 Subject: [PATCH 2/6] spacing --- src/latexify/codegen/algorithmic_codegen.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py index de951e3..5d33f45 100644 --- a/src/latexify/codegen/algorithmic_codegen.py +++ b/src/latexify/codegen/algorithmic_codegen.py @@ -64,7 +64,6 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str: ] latex = f"{self._prefix()}\\begin{{algorithmic}}\n" - self._indent += 1 latex += ( f"{self._prefix()}\\Function{{{node.name}}}{{${', '.join(arg_strs)}$}}\n" @@ -78,7 +77,6 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str: latex += f"{body_latex}\n{self._prefix()}\\EndFunction\n" self._indent -= 1 - return latex + rf"{self._prefix()}\end{{algorithmic}}" # TODO(ZibingZhang): support \ELSIF From c870af98d40c2bc53359d674e268dfe30cd60f7b Mon Sep 17 00:00:00 2001 From: Zibing Zhang Date: Mon, 12 Dec 2022 23:36:12 +0000 Subject: [PATCH 3/6] rm comment --- src/integration_tests/algorithmic_style_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/integration_tests/algorithmic_style_test.py b/src/integration_tests/algorithmic_style_test.py index 8ffdeaf..0c4f4db 100644 --- a/src/integration_tests/algorithmic_style_test.py +++ b/src/integration_tests/algorithmic_style_test.py @@ -35,7 +35,6 @@ def fact(n): else: return n * fact(n - 1) - # noqa: E501 latex = textwrap.dedent( r""" \begin{algorithmic} From a5a8ed7eef23aa03ea059a6a2ad99b1f07db1cca Mon Sep 17 00:00:00 2001 From: Zibing Zhang Date: Tue, 13 Dec 2022 08:08:32 +0000 Subject: [PATCH 4/6] prefix --- src/latexify/codegen/algorithmic_codegen.py | 78 +++++++++++++-------- 1 file changed, 47 insertions(+), 31 deletions(-) diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py index 5d33f45..929acfa 100644 --- a/src/latexify/codegen/algorithmic_codegen.py +++ b/src/latexify/codegen/algorithmic_codegen.py @@ -3,6 +3,8 @@ from __future__ import annotations import ast +import contextlib +from typing import Generator from latexify import exceptions from latexify.codegen import expression_codegen, identifier_converter @@ -18,7 +20,7 @@ class AlgorithmicCodegen(ast.NodeVisitor): _SPACES_PER_INDENT = 4 _identifier_converter: identifier_converter.IdentifierConverter - _indent: int + _indent_level: int def __init__( self, *, use_math_symbols: bool = False, use_set_symbols: bool = False @@ -36,7 +38,7 @@ def __init__( self._identifier_converter = identifier_converter.IdentifierConverter( use_math_symbols=use_math_symbols ) - self._indent = 0 + self._indent_level = 0 def generic_visit(self, node: ast.AST) -> str: raise exceptions.LatexifyNotSupportedError( @@ -50,11 +52,13 @@ def visit_Assign(self, node: ast.Assign) -> str: ] operands.append(self._expression_codegen.visit(node.value)) operands_latex = r" \gets ".join(operands) - return rf"{self._prefix()}\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"{self._prefix()}\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.""" @@ -63,39 +67,40 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str: self._identifier_converter.convert(arg.arg)[0] for arg in node.args.args ] - latex = f"{self._prefix()}\\begin{{algorithmic}}\n" - self._indent += 1 - latex += ( - f"{self._prefix()}\\Function{{{node.name}}}{{${', '.join(arg_strs)}$}}\n" + latex = self._add_indent("\\begin{algorithmic}\n") + self._indent_level += 1 + latex += self._add_indent( + f"\\Function{{{node.name}}}{{${', '.join(arg_strs)}$}}\n" ) # Body - self._indent += 1 + self._indent_level += 1 body_strs: list[str] = [self.visit(stmt) for stmt in node.body] - self._indent -= 1 + self._indent_level -= 1 body_latex = "\n".join(body_strs) - latex += f"{body_latex}\n{self._prefix()}\\EndFunction\n" - self._indent -= 1 - return latex + rf"{self._prefix()}\end{{algorithmic}}" + latex += f"{body_latex}\n" + latex += self._add_indent("\\EndFunction\n") + self._indent_level -= 1 + 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) - self._indent += 1 + self._indent_level += 1 body_latex = "\n".join(self.visit(stmt) for stmt in node.body) - self._indent -= 1 + self._indent_level -= 1 - latex = f"{self._prefix()}\\If{{${cond_latex}$}}\n{body_latex}" + latex = self._add_indent(f"\\If{{${cond_latex}$}}\n{body_latex}") if node.orelse: - latex += f"\n{self._prefix()}\\Else\n" - self._indent += 1 + latex += "\n" + self._add_indent(r"\Else") + "\n" + self._indent_level += 1 latex += "\n".join(self.visit(stmt) for stmt in node.orelse) - self._indent -= 1 + self._indent_level -= 1 - return latex + f"\n{self._prefix()}\\EndIf" + return latex + "\n" + self._add_indent(r"\EndIf") def visit_Module(self, node: ast.Module) -> str: """Visit a Module node.""" @@ -104,12 +109,11 @@ def visit_Module(self, node: ast.Module) -> str: def visit_Return(self, node: ast.Return) -> str: """Visit a Return node.""" return ( - ( - rf"{self._prefix()}\State \Return" - f" ${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 rf"{self._prefix()}\State \Return" + else self._add_indent(r"\State \Return") ) def visit_While(self, node: ast.While) -> str: @@ -120,14 +124,26 @@ def visit_While(self, node: ast.While) -> str: ) cond_latex = self._expression_codegen.visit(node.test) - self._indent += 1 + self._indent_level += 1 body_latex = "\n".join(self.visit(stmt) for stmt in node.body) - self._indent -= 1 + self._indent_level -= 1 return ( - f"{self._prefix()}\\While{{${cond_latex}$}}\n" - f"{body_latex}\n" - rf"{self._prefix()}\EndWhile" + self._add_indent(f"\\While{{${cond_latex}$}}\n") + + f"{body_latex}\n" + + self._add_indent(r"\EndWhile") ) - def _prefix(self) -> str: - return self._indent * self._SPACES_PER_INDENT * " " + @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 From 405831f3b3cf25abfec50c6356cfabcc5946a9ed Mon Sep 17 00:00:00 2001 From: Zibing Zhang Date: Tue, 13 Dec 2022 08:10:15 +0000 Subject: [PATCH 5/6] context mgr --- src/latexify/codegen/algorithmic_codegen.py | 37 +++++++++------------ 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py index 929acfa..9fd41f8 100644 --- a/src/latexify/codegen/algorithmic_codegen.py +++ b/src/latexify/codegen/algorithmic_codegen.py @@ -68,37 +68,33 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str: ] latex = self._add_indent("\\begin{algorithmic}\n") - self._indent_level += 1 - latex += self._add_indent( - f"\\Function{{{node.name}}}{{${', '.join(arg_strs)}$}}\n" - ) + with self._increment_level(): + latex += self._add_indent( + f"\\Function{{{node.name}}}{{${', '.join(arg_strs)}$}}\n" + ) - # Body - self._indent_level += 1 - body_strs: list[str] = [self.visit(stmt) for stmt in node.body] - self._indent_level -= 1 - body_latex = "\n".join(body_strs) + 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") - self._indent_level -= 1 + 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) - self._indent_level += 1 - body_latex = "\n".join(self.visit(stmt) for stmt in node.body) - self._indent_level -= 1 + with self._increment_level(): + body_latex = "\n".join(self.visit(stmt) for stmt in node.body) latex = self._add_indent(f"\\If{{${cond_latex}$}}\n{body_latex}") if node.orelse: latex += "\n" + self._add_indent(r"\Else") + "\n" - self._indent_level += 1 - latex += "\n".join(self.visit(stmt) for stmt in node.orelse) - self._indent_level -= 1 + with self._increment_level(): + latex += "\n".join(self.visit(stmt) for stmt in node.orelse) return latex + "\n" + self._add_indent(r"\EndIf") @@ -124,9 +120,8 @@ def visit_While(self, node: ast.While) -> str: ) cond_latex = self._expression_codegen.visit(node.test) - self._indent_level += 1 - body_latex = "\n".join(self.visit(stmt) for stmt in node.body) - self._indent_level -= 1 + 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" From 646af5b2a7fb04364bf555aa79e9376bb3bcc659 Mon Sep 17 00:00:00 2001 From: Zibing Zhang <44979059+ZibingZhang@users.noreply.github.com> Date: Tue, 13 Dec 2022 08:30:25 +0000 Subject: [PATCH 6/6] Update src/latexify/codegen/algorithmic_codegen.py Co-authored-by: Yusuke Oda --- src/latexify/codegen/algorithmic_codegen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py index 9fd41f8..0460ffd 100644 --- a/src/latexify/codegen/algorithmic_codegen.py +++ b/src/latexify/codegen/algorithmic_codegen.py @@ -4,7 +4,7 @@ import ast import contextlib -from typing import Generator +from collections.abc import Generator from latexify import exceptions from latexify.codegen import expression_codegen, identifier_converter