From 2f3d0608baebed0965cad4358293425799189f95 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 09:34:19 +0000
Subject: [PATCH 01/27] fe decorator

---
 src/latexify/frontend.py | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/src/latexify/frontend.py b/src/latexify/frontend.py
index 83a5d7e..eac330b 100644
--- a/src/latexify/frontend.py
+++ b/src/latexify/frontend.py
@@ -135,6 +135,43 @@ def _repr_latex_(self):
         )
 
 
+@overload
+def algorithmic(fn: Callable[..., Any], **kwargs: Any) -> LatexifiedFunction:
+    ...
+
+
+@overload
+def algorithmic(**kwargs: Any) -> LatexifiedFunction:
+    ...
+
+
+def algorithmic(
+    alg: Callable[..., Any] | None = None, **kwargs: Any
+) -> LatexifiedFunction | Callable[[Callable[..., Any]], LatexifiedFunction]:
+    """Attach LaTeX pretty-printing to the given algorithm.
+
+    This function works with or without specifying the target algorithm as the positional
+    argument. The following two syntaxes works similarly.
+        - latexify.algorithmic(alg, **kwargs)
+        - latexify.algorithmic(**kwargs)(alg)
+
+    Args:
+        alg: Callable to be wrapped.
+        **kwargs: Arguments to control behavior. See also get_latex().
+
+    Returns:
+        - If `alg` is passed, returns the wrapped function.
+        - Otherwise, returns the wrapper function with given settings.
+    """
+    if alg is not None:
+        return LatexifiedFunction(alg, **kwargs)
+
+    def wrapper(a):
+        return LatexifiedFunction(a, **kwargs)
+
+    return wrapper
+
+
 @overload
 def function(fn: Callable[..., Any], **kwargs: Any) -> LatexifiedFunction:
     ...

From 8d90da07572391d1b3a9cb46b64d5776e349f6f7 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 10:06:15 +0000
Subject: [PATCH 02/27] infra

---
 src/latexify/codegen/__init__.py            |   1 +
 src/latexify/codegen/algorithmic_codegen.py |  27 +++++
 src/latexify/frontend.py                    | 125 ++++++++------------
 src/latexify/output.py                      | 125 ++++++++++++++++++++
 4 files changed, 202 insertions(+), 76 deletions(-)
 create mode 100644 src/latexify/output.py

diff --git a/src/latexify/codegen/__init__.py b/src/latexify/codegen/__init__.py
index 1aea2c3..f028adf 100644
--- a/src/latexify/codegen/__init__.py
+++ b/src/latexify/codegen/__init__.py
@@ -3,5 +3,6 @@
 from latexify.codegen import algorithmic_codegen, expression_codegen, function_codegen
 
 AlgorithmicCodegen = algorithmic_codegen.AlgorithmicCodegen
+AlgorithmicJupyterCodegen = algorithmic_codegen.AlgorithmicJupyterCodegen
 ExpressionCodegen = expression_codegen.ExpressionCodegen
 FunctionCodegen = function_codegen.FunctionCodegen
diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index 2fcb16b..8b026e6 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -104,3 +104,30 @@ 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"
+
+
+class AlgorithmicJupyterCodegen(ast.NodeVisitor):
+    """Codegen for single algorithms, targeting the Jupyter Notebook environment.
+
+    This codegen works for Module with single FunctionDef node to generate a single
+    LaTeX expression of the given algorithm.
+    """
+
+    _identifier_converter: identifier_converter.IdentifierConverter
+
+    def __init__(
+        self, *, use_math_symbols: bool = False, use_set_symbols: bool = False
+    ) -> None:
+        """Initializer.
+
+        Args:
+            use_math_symbols: Whether to convert identifiers with a math symbol surface
+                (e.g., "alpha") to the LaTeX symbol (e.g., "\\alpha").
+            use_set_symbols: Whether to use set symbols or not.
+        """
+        self._expression_codegen = expression_codegen.ExpressionCodegen(
+            use_math_symbols=use_math_symbols, use_set_symbols=use_set_symbols
+        )
+        self._identifier_converter = identifier_converter.IdentifierConverter(
+            use_math_symbols=use_math_symbols
+        )
diff --git a/src/latexify/frontend.py b/src/latexify/frontend.py
index eac330b..f89a2a2 100644
--- a/src/latexify/frontend.py
+++ b/src/latexify/frontend.py
@@ -8,7 +8,7 @@
 
 from latexify import codegen
 from latexify import config as cfg
-from latexify import exceptions, parser, transformers
+from latexify import output, parser, transformers
 
 # NOTE(odashi):
 # These prefixes are trimmed by default.
@@ -17,15 +17,21 @@
 _COMMON_PREFIXES = {"math", "numpy", "np"}
 
 
+class Environment(enum.Enum):
+    JUPYTER_NOTEBOOK = "jupyter-notebook"
+    LATEX = "latex"
+
+
 class Style(enum.Enum):
+    ALGORITHMIC = "algorithmic"
     EXPRESSION = "expression"
     FUNCTION = "function"
-    ALGORITHMIC = "algorithmic"
 
 
 def get_latex(
     fn: Callable[..., Any],
     *,
+    environment: Environment = Environment.LATEX,
     style: Style = Style.FUNCTION,
     config: cfg.Config | None = None,
     **kwargs,
@@ -34,6 +40,7 @@ def get_latex(
 
     Args:
         fn: Reference to a function to analyze.
+        environment: Environment to target, the default is LATEX.
         style: Style of the LaTeX description, the default is FUNCTION.
         config: Use defined Config object, if it is None, it will be automatic assigned
             with default value.
@@ -68,10 +75,16 @@ def get_latex(
 
     # Generates LaTeX.
     if style == Style.ALGORITHMIC:
-        return codegen.AlgorithmicCodegen(
-            use_math_symbols=merged_config.use_math_symbols,
-            use_set_symbols=merged_config.use_set_symbols,
-        ).visit(tree)
+        if environment == Environment.LATEX:
+            return codegen.AlgorithmicCodegen(
+                use_math_symbols=merged_config.use_math_symbols,
+                use_set_symbols=merged_config.use_set_symbols,
+            ).visit(tree)
+        else:
+            return codegen.AlgorithmicJupyterCodegen(
+                use_math_symbols=merged_config.use_math_symbols,
+                use_set_symbols=merged_config.use_set_symbols,
+            ).visit(tree)
     else:
         return codegen.FunctionCodegen(
             use_math_symbols=merged_config.use_math_symbols,
@@ -80,78 +93,27 @@ def get_latex(
         ).visit(tree)
 
 
-class LatexifiedFunction:
-    """Function with latex representation."""
-
-    _fn: Callable[..., Any]
-    _latex: str | None
-    _error: str | None
-
-    def __init__(self, fn, **kwargs):
-        self._fn = fn
-        try:
-            self._latex = get_latex(fn, **kwargs)
-            self._error = None
-        except exceptions.LatexifyError as e:
-            self._latex = None
-            self._error = f"{type(e).__name__}: {str(e)}"
-
-    @property
-    def __doc__(self):
-        return self._fn.__doc__
-
-    @__doc__.setter
-    def __doc__(self, val):
-        self._fn.__doc__ = val
-
-    @property
-    def __name__(self):
-        return self._fn.__name__
-
-    @__name__.setter
-    def __name__(self, val):
-        self._fn.__name__ = val
-
-    def __call__(self, *args):
-        return self._fn(*args)
-
-    def __str__(self):
-        return self._latex if self._latex is not None else self._error
-
-    def _repr_html_(self):
-        """IPython hook to display HTML visualization."""
-        return (
-            '<span style="color: red;">' + self._error + "</span>"
-            if self._error is not None
-            else None
-        )
-
-    def _repr_latex_(self):
-        """IPython hook to display LaTeX visualization."""
-        return (
-            r"$$ \displaystyle " + self._latex + " $$"
-            if self._latex is not None
-            else self._error
-        )
-
-
 @overload
-def algorithmic(fn: Callable[..., Any], **kwargs: Any) -> LatexifiedFunction:
+def algorithmic(alg: Callable[..., Any], **kwargs: Any) -> output.LatexifiedAlgorithm:
     ...
 
 
 @overload
-def algorithmic(**kwargs: Any) -> LatexifiedFunction:
+def algorithmic(
+    **kwargs: Any,
+) -> Callable[[Callable[..., Any]], output.LatexifiedAlgorithm]:
     ...
 
 
 def algorithmic(
     alg: Callable[..., Any] | None = None, **kwargs: Any
-) -> LatexifiedFunction | Callable[[Callable[..., Any]], LatexifiedFunction]:
+) -> output.LatexifiedAlgorithm | Callable[
+    [Callable[..., Any]], output.LatexifiedAlgorithm
+]:
     """Attach LaTeX pretty-printing to the given algorithm.
 
-    This function works with or without specifying the target algorithm as the positional
-    argument. The following two syntaxes works similarly.
+    This function works with or without specifying the target algorithm as the
+    positional argument. The following two syntaxes works similarly.
         - latexify.algorithmic(alg, **kwargs)
         - latexify.algorithmic(**kwargs)(alg)
 
@@ -163,28 +125,35 @@ def algorithmic(
         - If `alg` is passed, returns the wrapped function.
         - Otherwise, returns the wrapper function with given settings.
     """
+    if "style" not in kwargs:
+        kwargs["style"] = Style.ALGORITHMIC
+
     if alg is not None:
-        return LatexifiedFunction(alg, **kwargs)
+        return output.LatexifiedAlgorithm(alg, **kwargs)
 
     def wrapper(a):
-        return LatexifiedFunction(a, **kwargs)
+        return output.LatexifiedAlgorithm(a, **kwargs)
 
     return wrapper
 
 
 @overload
-def function(fn: Callable[..., Any], **kwargs: Any) -> LatexifiedFunction:
+def function(fn: Callable[..., Any], **kwargs: Any) -> output.LatexifiedFunction:
     ...
 
 
 @overload
-def function(**kwargs: Any) -> Callable[[Callable[..., Any]], LatexifiedFunction]:
+def function(
+    **kwargs: Any,
+) -> Callable[[Callable[..., Any]], output.LatexifiedFunction]:
     ...
 
 
 def function(
     fn: Callable[..., Any] | None = None, **kwargs: Any
-) -> LatexifiedFunction | Callable[[Callable[..., Any]], LatexifiedFunction]:
+) -> output.LatexifiedFunction | Callable[
+    [Callable[..., Any]], output.LatexifiedFunction
+]:
     """Attach LaTeX pretty-printing to the given function.
 
     This function works with or without specifying the target function as the positional
@@ -201,27 +170,31 @@ def function(
         - Otherwise, returns the wrapper function with given settings.
     """
     if fn is not None:
-        return LatexifiedFunction(fn, **kwargs)
+        return output.LatexifiedFunction(fn, **kwargs)
 
     def wrapper(f):
-        return LatexifiedFunction(f, **kwargs)
+        return output.LatexifiedFunction(f, **kwargs)
 
     return wrapper
 
 
 @overload
-def expression(fn: Callable[..., Any], **kwargs: Any) -> LatexifiedFunction:
+def expression(fn: Callable[..., Any], **kwargs: Any) -> output.LatexifiedFunction:
     ...
 
 
 @overload
-def expression(**kwargs: Any) -> Callable[[Callable[..., Any]], LatexifiedFunction]:
+def expression(
+    **kwargs: Any,
+) -> Callable[[Callable[..., Any]], output.LatexifiedFunction]:
     ...
 
 
 def expression(
     fn: Callable[..., Any] | None = None, **kwargs: Any
-) -> LatexifiedFunction | Callable[[Callable[..., Any]], LatexifiedFunction]:
+) -> output.LatexifiedFunction | Callable[
+    [Callable[..., Any]], output.LatexifiedFunction
+]:
     """Attach LaTeX pretty-printing to the given function.
 
     This function is a shortcut for `latexify.function` with the default parameter
diff --git a/src/latexify/output.py b/src/latexify/output.py
new file mode 100644
index 0000000..bf77ecf
--- /dev/null
+++ b/src/latexify/output.py
@@ -0,0 +1,125 @@
+"""Output of the frontend function decorators."""
+
+from __future__ import annotations
+
+import abc
+import functools
+from typing import Any, Callable
+
+from latexify import exceptions, frontend
+
+
+class LatexifiedRepr(abc.ABC):
+    """Object with LaTeX representation."""
+
+    _fn: Callable[..., Any]
+    _latex: str | None
+    _error: str | None
+
+    def __init__(self, fn, **kwargs):
+        self._fn = fn
+
+    @property
+    def __doc__(self):
+        return self._fn.__doc__
+
+    @__doc__.setter
+    def __doc__(self, val):
+        self._fn.__doc__ = val
+
+    @property
+    def __name__(self):
+        return self._fn.__name__
+
+    @__name__.setter
+    def __name__(self, val):
+        self._fn.__name__ = val
+
+    def __call__(self, *args):
+        return self._fn(*args)
+
+    def __str__(self):
+        return self._latex if self._latex is not None else self._error
+
+    @abc.abstractmethod
+    def _repr_html_(self):
+        """IPython hook to display HTML visualization."""
+        ...
+
+    @abc.abstractmethod
+    def _repr_latex_(self):
+        """IPython hook to display LaTeX visualization."""
+        ...
+
+
+class LatexifiedAlgorithm(LatexifiedRepr):
+    """Algorithm with latex representation."""
+
+    _latex: str | None
+    _error: str | None
+    _jupyter_latex: str | None
+    _jupyter_error: str | None
+
+    def __init__(self, fn, **kwargs):
+        super().__init__(fn)
+        try:
+            kwargs["environment"] = frontend.Environment.LATEX
+            self._latex = frontend.get_latex(fn, **kwargs)
+            self._error = None
+        except exceptions.LatexifyError as e:
+            self._latex = None
+            self._error = f"{type(e).__name__}: {str(e)}"
+        try:
+            kwargs["environment"] = frontend.Environment.JUPYTER_NOTEBOOK
+            self._jupyter_latex = frontend.get_latex(fn, **kwargs)
+            self._jupyter_error = None
+        except exceptions.LatexifyError as e:
+            self._jupyter_latex = None
+            self._jupyter_error = f"{type(e).__name__}: {str(e)}"
+
+    @functools.cached_property
+    def _repr_html_(self):
+        """IPython hook to display HTML visualization."""
+        return (
+            '<span style="color: red;">' + self._jupyter_error + "</span>"
+            if self._jupyter_error is not None
+            else None
+        )
+
+    @functools.cached_property
+    def _repr_latex_(self):
+        """IPython hook to display LaTeX visualization."""
+        return (
+            r"$$ \displaystyle " + self._jupyter_latex + " $$"
+            if self._jupyter_latex is not None
+            else self._jupyter_error
+        )
+
+
+class LatexifiedFunction(LatexifiedRepr):
+    """Function with latex representation."""
+
+    def __init__(self, fn, **kwargs):
+        super().__init__(fn, **kwargs)
+        try:
+            self._latex = frontend.get_latex(fn, **kwargs)
+            self._error = None
+        except exceptions.LatexifyError as e:
+            self._latex = None
+            self._error = f"{type(e).__name__}: {str(e)}"
+
+    def _repr_html_(self):
+        """IPython hook to display HTML visualization."""
+        return (
+            '<span style="color: red;">' + self._error + "</span>"
+            if self._error is not None
+            else None
+        )
+
+    def _repr_latex_(self):
+        """IPython hook to display LaTeX visualization."""
+        return (
+            r"$$ \displaystyle " + self._latex + " $$"
+            if self._latex is not None
+            else self._error
+        )

From 200b7fc91c459ad1a45c25ebdb1bddcdd6c49af1 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 10:12:58 +0000
Subject: [PATCH 03/27] no longer cached

---
 src/latexify/output.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/latexify/output.py b/src/latexify/output.py
index bf77ecf..26d433f 100644
--- a/src/latexify/output.py
+++ b/src/latexify/output.py
@@ -77,7 +77,6 @@ def __init__(self, fn, **kwargs):
             self._jupyter_latex = None
             self._jupyter_error = f"{type(e).__name__}: {str(e)}"
 
-    @functools.cached_property
     def _repr_html_(self):
         """IPython hook to display HTML visualization."""
         return (
@@ -86,11 +85,10 @@ def _repr_html_(self):
             else None
         )
 
-    @functools.cached_property
     def _repr_latex_(self):
         """IPython hook to display LaTeX visualization."""
         return (
-            r"$$ \displaystyle " + self._jupyter_latex + " $$"
+            r"$ " + self._jupyter_latex + " $"
             if self._jupyter_latex is not None
             else self._jupyter_error
         )

From d71392c70d131913b6ef64653ba583110d7e1d11 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 10:13:06 +0000
Subject: [PATCH 04/27] rm import

---
 src/latexify/output.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/latexify/output.py b/src/latexify/output.py
index 26d433f..662cdb2 100644
--- a/src/latexify/output.py
+++ b/src/latexify/output.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import abc
-import functools
 from typing import Any, Callable
 
 from latexify import exceptions, frontend

From a2bf5248544f2c2705d8ffbb7ff20d98cfb21948 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 11:34:31 +0000
Subject: [PATCH 05/27] codegen

---
 .../algorithmic_style_test.py                 | 50 ++++++++++
 src/latexify/codegen/algorithmic_codegen.py   | 92 ++++++++++++++++++-
 2 files changed, 141 insertions(+), 1 deletion(-)

diff --git a/src/integration_tests/algorithmic_style_test.py b/src/integration_tests/algorithmic_style_test.py
index 047d6f9..143b3e1 100644
--- a/src/integration_tests/algorithmic_style_test.py
+++ b/src/integration_tests/algorithmic_style_test.py
@@ -48,6 +48,27 @@ def fact(n):
     check_algorithm(fact, latex)
 
 
+def test_factorial_jupyter() -> None:
+    def fact(n):
+        if n == 0:
+            return 1
+        else:
+            return n * fact(n - 1)
+
+    latex = (
+        r"\displaystyle \hspace{0pt} \mathbf{function} \ \mathrm{FACT}(n): \\"
+        r" \displaystyle \hspace{16pt}\mathbf{if} \ n = 0 \\"
+        r" \displaystyle \hspace{32pt} \mathbf{return} \ 1 \\"
+        r" \displaystyle \hspace{16pt}\mathbf{else} \\"
+        r" \displaystyle \hspace{32pt}"
+        r" \mathbf{return} \ n"
+        r" \mathrm{fact} \mathopen{}\left( n - 1 \mathclose{}\right) \\"
+        r" \displaystyle \hspace{16pt} \mathbf{end \ if} \\"
+        r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+    )
+    check_algorithm(fact, latex, environment=frontend.Environment.JUPYTER_NOTEBOOK)
+
+
 def test_collatz() -> None:
     def collatz(n):
         iterations = 0
@@ -75,3 +96,32 @@ def collatz(n):
         r" \end{algorithmic}"
     )
     check_algorithm(collatz, latex)
+
+
+def test_collatz_jupyter() -> None:
+    def collatz(n):
+        iterations = 0
+        while n > 1:
+            if n % 2 == 0:
+                n = n // 2
+            else:
+                n = 3 * n + 1
+            iterations = iterations + 1
+        return iterations
+
+    latex = (
+        r"\displaystyle \hspace{0pt} \mathbf{function} \ \mathrm{COLLATZ}(n): \\"
+        r" \displaystyle \hspace{16pt} \mathrm{iterations} \gets 0 \\"
+        r" \displaystyle \hspace{16pt} \mathbf{while} \ n > 1 \\"
+        r" \displaystyle \hspace{32pt}\mathbf{if} \ n \mathbin{\%} 2 = 0 \\"
+        r" \displaystyle \hspace{48pt} n \gets \left\lfloor\frac{n}{2}\right\rfloor \\"
+        r" \displaystyle \hspace{32pt}\mathbf{else} \\"
+        r" \displaystyle \hspace{48pt} n \gets 3 n + 1 \\"
+        r" \displaystyle \hspace{32pt} \mathbf{end \ if} \\"
+        r" \displaystyle \hspace{32pt}"
+        r" \mathrm{iterations} \gets \mathrm{iterations} + 1 \\"
+        r" \displaystyle \hspace{16pt} \mathbf{end \ while} \\"
+        r" \displaystyle \hspace{16pt} \mathbf{return} \ \mathrm{iterations} \\"
+        r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+    )
+    check_algorithm(collatz, latex, environment=frontend.Environment.JUPYTER_NOTEBOOK)
diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index 8b026e6..b93d8e3 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -60,6 +60,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         ]
         # Body
         body_strs: list[str] = [self.visit(stmt) for stmt in node.body]
+
         return (
             rf"\begin{{algorithmic}}"
             rf" \Function{{{node.name}}}{{${', '.join(arg_strs)}$}}"
@@ -73,7 +74,6 @@ 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)
-
         latex = rf"\If{{${cond_latex}$}} {body_latex}"
 
         if node.orelse:
@@ -113,7 +113,9 @@ class AlgorithmicJupyterCodegen(ast.NodeVisitor):
     LaTeX expression of the given algorithm.
     """
 
+    _PT_PER_INDENT = 16
     _identifier_converter: identifier_converter.IdentifierConverter
+    _indent: int
 
     def __init__(
         self, *, use_math_symbols: bool = False, use_set_symbols: bool = False
@@ -131,3 +133,91 @@ 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(
+            f"Unsupported AST: {type(node).__name__}"
+        )
+
+    def visit_Assign(self, node: ast.Assign) -> str:
+        """Visit an Assign node."""
+        operands: list[str] = [
+            self._expression_codegen.visit(target) for target in node.targets
+        ]
+        operands.append(self._expression_codegen.visit(node.value))
+        operands_latex = r" \gets ".join(operands)
+        return rf"{self._prefix()} {operands_latex}"
+
+    def visit_Expr(self, node: ast.Expr) -> str:
+        """Visit an Expr node."""
+        return rf"{self._prefix()} {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
+        self._indent += 1
+        body_strs: list[str] = [self.visit(stmt) for stmt in node.body]
+        self._indent -= 1
+        body = r" \\ ".join(body_strs)
+
+        return (
+            rf"{self._prefix()} \mathbf{{function}}"
+            rf" \ \mathrm{{{node.name.upper()}}}({', '.join(arg_strs)}): \\"
+            rf" {body} \\"
+            rf" {self._prefix()} \mathbf{{end \ function}}"
+        )
+
+    # 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
+        body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
+        self._indent -= 1
+        latex = rf"{self._prefix()}\mathbf{{if}} \ {cond_latex} \\ {body_latex}"
+
+        if node.orelse:
+            latex += rf" \\ {self._prefix()}\mathbf{{else}} \\ "
+            self._indent += 1
+            latex += r" \\ ".join(self.visit(stmt) for stmt in node.orelse)
+            self._indent -= 1
+
+        return latex + rf" \\ {self._prefix()} \mathbf{{end \ if}}"
+
+    def visit_Module(self, node: ast.Module) -> str:
+        """Visit a Module node."""
+        return self.visit(node.body[0])
+
+    def visit_Return(self, node: ast.Return) -> str:
+        """Visit a Return node."""
+        return (
+            rf"{self._prefix()} \mathbf{{return}}"
+            rf" \ {self._expression_codegen.visit(node.value)}"
+            if node.value is not None
+            else rf"{self._prefix()} \mathbf{{return}}"
+        )
+
+    def visit_While(self, node: ast.While) -> str:
+        """Visit a While node."""
+        if node.orelse:
+            raise exceptions.LatexifyNotSupportedError(
+                "While statement with the else clause is not supported"
+            )
+
+        cond_latex = self._expression_codegen.visit(node.test)
+        self._indent += 1
+        body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
+        self._indent -= 1
+        return (
+            rf"{self._prefix()} \mathbf{{while}} \ {cond_latex} \\ "
+            rf"{body_latex} \\ "
+            rf"{self._prefix()} \mathbf{{end \ while}}"
+        )
+
+    def _prefix(self) -> str:
+        return rf"\displaystyle \hspace{{{self._indent * self._PT_PER_INDENT}pt}}"

From 6d9411c5a9bf21b7d84b667c1e8c3910fd75c016 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 11:51:11 +0000
Subject: [PATCH 06/27] tests & bugs

---
 .../algorithmic_style_test.py                 |  48 +++----
 src/latexify/codegen/algorithmic_codegen.py   |   8 +-
 .../codegen/algorithmic_codegen_test.py       | 128 ++++++++++++++++++
 3 files changed, 156 insertions(+), 28 deletions(-)

diff --git a/src/integration_tests/algorithmic_style_test.py b/src/integration_tests/algorithmic_style_test.py
index 143b3e1..a6d6d7a 100644
--- a/src/integration_tests/algorithmic_style_test.py
+++ b/src/integration_tests/algorithmic_style_test.py
@@ -48,27 +48,6 @@ def fact(n):
     check_algorithm(fact, latex)
 
 
-def test_factorial_jupyter() -> None:
-    def fact(n):
-        if n == 0:
-            return 1
-        else:
-            return n * fact(n - 1)
-
-    latex = (
-        r"\displaystyle \hspace{0pt} \mathbf{function} \ \mathrm{FACT}(n): \\"
-        r" \displaystyle \hspace{16pt}\mathbf{if} \ n = 0 \\"
-        r" \displaystyle \hspace{32pt} \mathbf{return} \ 1 \\"
-        r" \displaystyle \hspace{16pt}\mathbf{else} \\"
-        r" \displaystyle \hspace{32pt}"
-        r" \mathbf{return} \ n"
-        r" \mathrm{fact} \mathopen{}\left( n - 1 \mathclose{}\right) \\"
-        r" \displaystyle \hspace{16pt} \mathbf{end \ if} \\"
-        r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
-    )
-    check_algorithm(fact, latex, environment=frontend.Environment.JUPYTER_NOTEBOOK)
-
-
 def test_collatz() -> None:
     def collatz(n):
         iterations = 0
@@ -98,6 +77,27 @@ def collatz(n):
     check_algorithm(collatz, latex)
 
 
+def test_factorial_jupyter() -> None:
+    def fact(n):
+        if n == 0:
+            return 1
+        else:
+            return n * fact(n - 1)
+
+    latex = (
+        r"\displaystyle \hspace{0pt} \mathbf{function} \ \mathrm{FACT}(n) \\"
+        r" \displaystyle \hspace{16pt} \mathbf{if} \ n = 0 \\"
+        r" \displaystyle \hspace{32pt} \mathbf{return} \ 1 \\"
+        r" \displaystyle \hspace{16pt} \mathbf{else} \\"
+        r" \displaystyle \hspace{32pt}"
+        r" \mathbf{return} \ n"
+        r" \mathrm{fact} \mathopen{}\left( n - 1 \mathclose{}\right) \\"
+        r" \displaystyle \hspace{16pt} \mathbf{end \ if} \\"
+        r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+    )
+    check_algorithm(fact, latex, environment=frontend.Environment.JUPYTER_NOTEBOOK)
+
+
 def test_collatz_jupyter() -> None:
     def collatz(n):
         iterations = 0
@@ -110,12 +110,12 @@ def collatz(n):
         return iterations
 
     latex = (
-        r"\displaystyle \hspace{0pt} \mathbf{function} \ \mathrm{COLLATZ}(n): \\"
+        r"\displaystyle \hspace{0pt} \mathbf{function} \ \mathrm{COLLATZ}(n) \\"
         r" \displaystyle \hspace{16pt} \mathrm{iterations} \gets 0 \\"
         r" \displaystyle \hspace{16pt} \mathbf{while} \ n > 1 \\"
-        r" \displaystyle \hspace{32pt}\mathbf{if} \ n \mathbin{\%} 2 = 0 \\"
+        r" \displaystyle \hspace{32pt} \mathbf{if} \ n \mathbin{\%} 2 = 0 \\"
         r" \displaystyle \hspace{48pt} n \gets \left\lfloor\frac{n}{2}\right\rfloor \\"
-        r" \displaystyle \hspace{32pt}\mathbf{else} \\"
+        r" \displaystyle \hspace{32pt} \mathbf{else} \\"
         r" \displaystyle \hspace{48pt} n \gets 3 n + 1 \\"
         r" \displaystyle \hspace{32pt} \mathbf{end \ if} \\"
         r" \displaystyle \hspace{32pt}"
diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index b93d8e3..973e9df 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -107,7 +107,7 @@ def visit_While(self, node: ast.While) -> str:
 
 
 class AlgorithmicJupyterCodegen(ast.NodeVisitor):
-    """Codegen for single algorithms, targeting the Jupyter Notebook environment.
+    """Codegen for single algorithms targeting the Jupyter Notebook environment.
 
     This codegen works for Module with single FunctionDef node to generate a single
     LaTeX expression of the given algorithm.
@@ -167,7 +167,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
 
         return (
             rf"{self._prefix()} \mathbf{{function}}"
-            rf" \ \mathrm{{{node.name.upper()}}}({', '.join(arg_strs)}): \\"
+            rf" \ \mathrm{{{node.name.upper()}}}({', '.join(arg_strs)}) \\"
             rf" {body} \\"
             rf" {self._prefix()} \mathbf{{end \ function}}"
         )
@@ -179,10 +179,10 @@ def visit_If(self, node: ast.If) -> str:
         self._indent += 1
         body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
         self._indent -= 1
-        latex = rf"{self._prefix()}\mathbf{{if}} \ {cond_latex} \\ {body_latex}"
+        latex = rf"{self._prefix()} \mathbf{{if}} \ {cond_latex} \\ {body_latex}"
 
         if node.orelse:
-            latex += rf" \\ {self._prefix()}\mathbf{{else}} \\ "
+            latex += rf" \\ {self._prefix()} \mathbf{{else}} \\ "
             self._indent += 1
             latex += r" \\ ".join(self.visit(stmt) for stmt in node.orelse)
             self._indent -= 1
diff --git a/src/latexify/codegen/algorithmic_codegen_test.py b/src/latexify/codegen/algorithmic_codegen_test.py
index a972beb..5051daf 100644
--- a/src/latexify/codegen/algorithmic_codegen_test.py
+++ b/src/latexify/codegen/algorithmic_codegen_test.py
@@ -133,3 +133,131 @@ def test_visit_while_with_else() -> None:
         match="^While statement with the else clause is not supported$",
     ):
         algorithmic_codegen.AlgorithmicCodegen().visit(node)
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        ("x = 3", r"\displaystyle \hspace{0pt} x \gets 3"),
+        ("a = b = 0", r"\displaystyle \hspace{0pt} a \gets b \gets 0"),
+    ],
+)
+def test_visit_assign_jupyter(code: str, latex: str) -> None:
+    node = ast.parse(textwrap.dedent(code)).body[0]
+    assert isinstance(node, ast.Assign)
+    assert algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        (
+            "def f(x): return x",
+            (
+                r"\displaystyle \hspace{0pt} \mathbf{function}"
+                r" \ \mathrm{F}(x) \\"
+                r" \displaystyle \hspace{16pt} \mathbf{return} \ x \\"
+                r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+            ),
+        ),
+        (
+            "def f(a, b, c): return 3",
+            (
+                r"\displaystyle \hspace{0pt} \mathbf{function}"
+                r" \ \mathrm{F}(a, b, c) \\"
+                r" \displaystyle \hspace{16pt} \mathbf{return} \ 3 \\"
+                r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+            ),
+        ),
+    ],
+)
+def test_visit_functiondef_jupyter(code: str, latex: str) -> None:
+    node = ast.parse(textwrap.dedent(code)).body[0]
+    assert isinstance(node, ast.FunctionDef)
+    assert algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        (
+            "if x < y: return x",
+            (
+                r"\displaystyle \hspace{0pt} \mathbf{if} \ x < y \\"
+                r" \displaystyle \hspace{16pt} \mathbf{return} \ x \\"
+                r" \displaystyle \hspace{0pt} \mathbf{end \ if}"
+            ),
+        ),
+        (
+            "if True: x\nelse: y",
+            (
+                r"\displaystyle \hspace{0pt} \mathbf{if} \ \mathrm{True} \\"
+                r" \displaystyle \hspace{16pt} x \\"
+                r" \displaystyle \hspace{0pt} \mathbf{else} \\"
+                r" \displaystyle \hspace{16pt} y \\"
+                r" \displaystyle \hspace{0pt} \mathbf{end \ if}"
+            ),
+        ),
+    ],
+)
+def test_visit_if_jupyter(code: str, latex: str) -> None:
+    node = ast.parse(textwrap.dedent(code)).body[0]
+    assert isinstance(node, ast.If)
+    assert algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        (
+            "return x + y",
+            r"\displaystyle \hspace{0pt} \mathbf{return} \ x + y",
+        ),
+        (
+            "return",
+            r"\displaystyle \hspace{0pt} \mathbf{return}",
+        ),
+    ],
+)
+def test_visit_return_jupyter(code: str, latex: str) -> None:
+    node = ast.parse(textwrap.dedent(code)).body[0]
+    assert isinstance(node, ast.Return)
+    assert algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        (
+            "while x < y: x = x + 1",
+            (
+                r"\displaystyle \hspace{0pt} \mathbf{while} \ x < y \\"
+                r" \displaystyle \hspace{16pt} x \gets x + 1 \\"
+                r" \displaystyle \hspace{0pt} \mathbf{end \ while}"
+            ),
+        )
+    ],
+)
+def test_visit_while_jupyter(code: str, latex: str) -> None:
+    node = ast.parse(textwrap.dedent(code)).body[0]
+    assert isinstance(node, ast.While)
+    assert algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node) == latex
+
+
+def test_visit_while_with_else_jupyter() -> None:
+    node = ast.parse(
+        textwrap.dedent(
+            """
+            while True:
+                x = x
+            else:
+                x = y
+            """
+        )
+    ).body[0]
+    assert isinstance(node, ast.While)
+    with pytest.raises(
+        exceptions.LatexifyNotSupportedError,
+        match="^While statement with the else clause is not supported$",
+    ):
+        algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node)

From f5dc568c65c47b63794cb5722698fe0b69cb4498 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 11:54:38 +0000
Subject: [PATCH 07/27] temp del

---
 src/latexify/output.py | 122 -----------------------------------------
 1 file changed, 122 deletions(-)
 delete mode 100644 src/latexify/output.py

diff --git a/src/latexify/output.py b/src/latexify/output.py
deleted file mode 100644
index 662cdb2..0000000
--- a/src/latexify/output.py
+++ /dev/null
@@ -1,122 +0,0 @@
-"""Output of the frontend function decorators."""
-
-from __future__ import annotations
-
-import abc
-from typing import Any, Callable
-
-from latexify import exceptions, frontend
-
-
-class LatexifiedRepr(abc.ABC):
-    """Object with LaTeX representation."""
-
-    _fn: Callable[..., Any]
-    _latex: str | None
-    _error: str | None
-
-    def __init__(self, fn, **kwargs):
-        self._fn = fn
-
-    @property
-    def __doc__(self):
-        return self._fn.__doc__
-
-    @__doc__.setter
-    def __doc__(self, val):
-        self._fn.__doc__ = val
-
-    @property
-    def __name__(self):
-        return self._fn.__name__
-
-    @__name__.setter
-    def __name__(self, val):
-        self._fn.__name__ = val
-
-    def __call__(self, *args):
-        return self._fn(*args)
-
-    def __str__(self):
-        return self._latex if self._latex is not None else self._error
-
-    @abc.abstractmethod
-    def _repr_html_(self):
-        """IPython hook to display HTML visualization."""
-        ...
-
-    @abc.abstractmethod
-    def _repr_latex_(self):
-        """IPython hook to display LaTeX visualization."""
-        ...
-
-
-class LatexifiedAlgorithm(LatexifiedRepr):
-    """Algorithm with latex representation."""
-
-    _latex: str | None
-    _error: str | None
-    _jupyter_latex: str | None
-    _jupyter_error: str | None
-
-    def __init__(self, fn, **kwargs):
-        super().__init__(fn)
-        try:
-            kwargs["environment"] = frontend.Environment.LATEX
-            self._latex = frontend.get_latex(fn, **kwargs)
-            self._error = None
-        except exceptions.LatexifyError as e:
-            self._latex = None
-            self._error = f"{type(e).__name__}: {str(e)}"
-        try:
-            kwargs["environment"] = frontend.Environment.JUPYTER_NOTEBOOK
-            self._jupyter_latex = frontend.get_latex(fn, **kwargs)
-            self._jupyter_error = None
-        except exceptions.LatexifyError as e:
-            self._jupyter_latex = None
-            self._jupyter_error = f"{type(e).__name__}: {str(e)}"
-
-    def _repr_html_(self):
-        """IPython hook to display HTML visualization."""
-        return (
-            '<span style="color: red;">' + self._jupyter_error + "</span>"
-            if self._jupyter_error is not None
-            else None
-        )
-
-    def _repr_latex_(self):
-        """IPython hook to display LaTeX visualization."""
-        return (
-            r"$ " + self._jupyter_latex + " $"
-            if self._jupyter_latex is not None
-            else self._jupyter_error
-        )
-
-
-class LatexifiedFunction(LatexifiedRepr):
-    """Function with latex representation."""
-
-    def __init__(self, fn, **kwargs):
-        super().__init__(fn, **kwargs)
-        try:
-            self._latex = frontend.get_latex(fn, **kwargs)
-            self._error = None
-        except exceptions.LatexifyError as e:
-            self._latex = None
-            self._error = f"{type(e).__name__}: {str(e)}"
-
-    def _repr_html_(self):
-        """IPython hook to display HTML visualization."""
-        return (
-            '<span style="color: red;">' + self._error + "</span>"
-            if self._error is not None
-            else None
-        )
-
-    def _repr_latex_(self):
-        """IPython hook to display LaTeX visualization."""
-        return (
-            r"$$ \displaystyle " + self._latex + " $$"
-            if self._latex is not None
-            else self._error
-        )

From b2d8f4c4ff2eda2348a420d0fb3f557c69e9e821 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 11:56:08 +0000
Subject: [PATCH 08/27] mv

---
 src/latexify/{frontend.py => output.py} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename src/latexify/{frontend.py => output.py} (100%)

diff --git a/src/latexify/frontend.py b/src/latexify/output.py
similarity index 100%
rename from src/latexify/frontend.py
rename to src/latexify/output.py

From a5419441a0985352ec9d0ba2bc1b5e21d4a9cac1 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 11:56:45 +0000
Subject: [PATCH 09/27] copy back

---
 src/latexify/frontend.py | 207 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 207 insertions(+)
 create mode 100644 src/latexify/frontend.py

diff --git a/src/latexify/frontend.py b/src/latexify/frontend.py
new file mode 100644
index 0000000..f89a2a2
--- /dev/null
+++ b/src/latexify/frontend.py
@@ -0,0 +1,207 @@
+"""Frontend interfaces of latexify."""
+
+from __future__ import annotations
+
+import enum
+from collections.abc import Callable
+from typing import Any, overload
+
+from latexify import codegen
+from latexify import config as cfg
+from latexify import output, parser, transformers
+
+# NOTE(odashi):
+# These prefixes are trimmed by default.
+# This behavior shouldn't be controlled by users in the current implementation because
+# some processes expects absense of these prefixes.
+_COMMON_PREFIXES = {"math", "numpy", "np"}
+
+
+class Environment(enum.Enum):
+    JUPYTER_NOTEBOOK = "jupyter-notebook"
+    LATEX = "latex"
+
+
+class Style(enum.Enum):
+    ALGORITHMIC = "algorithmic"
+    EXPRESSION = "expression"
+    FUNCTION = "function"
+
+
+def get_latex(
+    fn: Callable[..., Any],
+    *,
+    environment: Environment = Environment.LATEX,
+    style: Style = Style.FUNCTION,
+    config: cfg.Config | None = None,
+    **kwargs,
+) -> str:
+    """Obtains LaTeX description from the function's source.
+
+    Args:
+        fn: Reference to a function to analyze.
+        environment: Environment to target, the default is LATEX.
+        style: Style of the LaTeX description, the default is FUNCTION.
+        config: Use defined Config object, if it is None, it will be automatic assigned
+            with default value.
+        **kwargs: Dict of Config field values that could be defined individually
+            by users.
+
+    Returns:
+        Generated LaTeX description.
+
+    Raises:
+        latexify.exceptions.LatexifyError: Something went wrong during conversion.
+    """
+    if style == Style.EXPRESSION:
+        kwargs["use_signature"] = kwargs.get("use_signature", False)
+
+    merged_config = cfg.Config.defaults().merge(config=config, **kwargs)
+
+    # Obtains the source AST.
+    tree = parser.parse_function(fn)
+
+    # Applies AST transformations.
+
+    prefixes = _COMMON_PREFIXES | (merged_config.prefixes or set())
+    tree = transformers.PrefixTrimmer(prefixes).visit(tree)
+
+    if merged_config.identifiers is not None:
+        tree = transformers.IdentifierReplacer(merged_config.identifiers).visit(tree)
+    if merged_config.reduce_assignments:
+        tree = transformers.AssignmentReducer().visit(tree)
+    if merged_config.expand_functions is not None:
+        tree = transformers.FunctionExpander(merged_config.expand_functions).visit(tree)
+
+    # Generates LaTeX.
+    if style == Style.ALGORITHMIC:
+        if environment == Environment.LATEX:
+            return codegen.AlgorithmicCodegen(
+                use_math_symbols=merged_config.use_math_symbols,
+                use_set_symbols=merged_config.use_set_symbols,
+            ).visit(tree)
+        else:
+            return codegen.AlgorithmicJupyterCodegen(
+                use_math_symbols=merged_config.use_math_symbols,
+                use_set_symbols=merged_config.use_set_symbols,
+            ).visit(tree)
+    else:
+        return codegen.FunctionCodegen(
+            use_math_symbols=merged_config.use_math_symbols,
+            use_signature=merged_config.use_signature,
+            use_set_symbols=merged_config.use_set_symbols,
+        ).visit(tree)
+
+
+@overload
+def algorithmic(alg: Callable[..., Any], **kwargs: Any) -> output.LatexifiedAlgorithm:
+    ...
+
+
+@overload
+def algorithmic(
+    **kwargs: Any,
+) -> Callable[[Callable[..., Any]], output.LatexifiedAlgorithm]:
+    ...
+
+
+def algorithmic(
+    alg: Callable[..., Any] | None = None, **kwargs: Any
+) -> output.LatexifiedAlgorithm | Callable[
+    [Callable[..., Any]], output.LatexifiedAlgorithm
+]:
+    """Attach LaTeX pretty-printing to the given algorithm.
+
+    This function works with or without specifying the target algorithm as the
+    positional argument. The following two syntaxes works similarly.
+        - latexify.algorithmic(alg, **kwargs)
+        - latexify.algorithmic(**kwargs)(alg)
+
+    Args:
+        alg: Callable to be wrapped.
+        **kwargs: Arguments to control behavior. See also get_latex().
+
+    Returns:
+        - If `alg` is passed, returns the wrapped function.
+        - Otherwise, returns the wrapper function with given settings.
+    """
+    if "style" not in kwargs:
+        kwargs["style"] = Style.ALGORITHMIC
+
+    if alg is not None:
+        return output.LatexifiedAlgorithm(alg, **kwargs)
+
+    def wrapper(a):
+        return output.LatexifiedAlgorithm(a, **kwargs)
+
+    return wrapper
+
+
+@overload
+def function(fn: Callable[..., Any], **kwargs: Any) -> output.LatexifiedFunction:
+    ...
+
+
+@overload
+def function(
+    **kwargs: Any,
+) -> Callable[[Callable[..., Any]], output.LatexifiedFunction]:
+    ...
+
+
+def function(
+    fn: Callable[..., Any] | None = None, **kwargs: Any
+) -> output.LatexifiedFunction | Callable[
+    [Callable[..., Any]], output.LatexifiedFunction
+]:
+    """Attach LaTeX pretty-printing to the given function.
+
+    This function works with or without specifying the target function as the positional
+    argument. The following two syntaxes works similarly.
+        - latexify.function(fn, **kwargs)
+        - latexify.function(**kwargs)(fn)
+
+    Args:
+        fn: Callable to be wrapped.
+        **kwargs: Arguments to control behavior. See also get_latex().
+
+    Returns:
+        - If `fn` is passed, returns the wrapped function.
+        - Otherwise, returns the wrapper function with given settings.
+    """
+    if fn is not None:
+        return output.LatexifiedFunction(fn, **kwargs)
+
+    def wrapper(f):
+        return output.LatexifiedFunction(f, **kwargs)
+
+    return wrapper
+
+
+@overload
+def expression(fn: Callable[..., Any], **kwargs: Any) -> output.LatexifiedFunction:
+    ...
+
+
+@overload
+def expression(
+    **kwargs: Any,
+) -> Callable[[Callable[..., Any]], output.LatexifiedFunction]:
+    ...
+
+
+def expression(
+    fn: Callable[..., Any] | None = None, **kwargs: Any
+) -> output.LatexifiedFunction | Callable[
+    [Callable[..., Any]], output.LatexifiedFunction
+]:
+    """Attach LaTeX pretty-printing to the given function.
+
+    This function is a shortcut for `latexify.function` with the default parameter
+    `use_signature=False`.
+    """
+    kwargs["style"] = Style.EXPRESSION
+    if fn is not None:
+        return function(fn, **kwargs)
+    else:
+        return function(**kwargs)

From 8296cd1036521481e2ef5d1d739b1c247e3eefbc Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 11:57:30 +0000
Subject: [PATCH 10/27] make edit

---
 src/latexify/output.py | 323 +++++++++++++++--------------------------
 1 file changed, 119 insertions(+), 204 deletions(-)

diff --git a/src/latexify/output.py b/src/latexify/output.py
index f89a2a2..662cdb2 100644
--- a/src/latexify/output.py
+++ b/src/latexify/output.py
@@ -1,207 +1,122 @@
-"""Frontend interfaces of latexify."""
+"""Output of the frontend function decorators."""
 
 from __future__ import annotations
 
-import enum
-from collections.abc import Callable
-from typing import Any, overload
-
-from latexify import codegen
-from latexify import config as cfg
-from latexify import output, parser, transformers
-
-# NOTE(odashi):
-# These prefixes are trimmed by default.
-# This behavior shouldn't be controlled by users in the current implementation because
-# some processes expects absense of these prefixes.
-_COMMON_PREFIXES = {"math", "numpy", "np"}
-
-
-class Environment(enum.Enum):
-    JUPYTER_NOTEBOOK = "jupyter-notebook"
-    LATEX = "latex"
-
-
-class Style(enum.Enum):
-    ALGORITHMIC = "algorithmic"
-    EXPRESSION = "expression"
-    FUNCTION = "function"
-
-
-def get_latex(
-    fn: Callable[..., Any],
-    *,
-    environment: Environment = Environment.LATEX,
-    style: Style = Style.FUNCTION,
-    config: cfg.Config | None = None,
-    **kwargs,
-) -> str:
-    """Obtains LaTeX description from the function's source.
-
-    Args:
-        fn: Reference to a function to analyze.
-        environment: Environment to target, the default is LATEX.
-        style: Style of the LaTeX description, the default is FUNCTION.
-        config: Use defined Config object, if it is None, it will be automatic assigned
-            with default value.
-        **kwargs: Dict of Config field values that could be defined individually
-            by users.
-
-    Returns:
-        Generated LaTeX description.
-
-    Raises:
-        latexify.exceptions.LatexifyError: Something went wrong during conversion.
-    """
-    if style == Style.EXPRESSION:
-        kwargs["use_signature"] = kwargs.get("use_signature", False)
-
-    merged_config = cfg.Config.defaults().merge(config=config, **kwargs)
-
-    # Obtains the source AST.
-    tree = parser.parse_function(fn)
-
-    # Applies AST transformations.
-
-    prefixes = _COMMON_PREFIXES | (merged_config.prefixes or set())
-    tree = transformers.PrefixTrimmer(prefixes).visit(tree)
-
-    if merged_config.identifiers is not None:
-        tree = transformers.IdentifierReplacer(merged_config.identifiers).visit(tree)
-    if merged_config.reduce_assignments:
-        tree = transformers.AssignmentReducer().visit(tree)
-    if merged_config.expand_functions is not None:
-        tree = transformers.FunctionExpander(merged_config.expand_functions).visit(tree)
-
-    # Generates LaTeX.
-    if style == Style.ALGORITHMIC:
-        if environment == Environment.LATEX:
-            return codegen.AlgorithmicCodegen(
-                use_math_symbols=merged_config.use_math_symbols,
-                use_set_symbols=merged_config.use_set_symbols,
-            ).visit(tree)
-        else:
-            return codegen.AlgorithmicJupyterCodegen(
-                use_math_symbols=merged_config.use_math_symbols,
-                use_set_symbols=merged_config.use_set_symbols,
-            ).visit(tree)
-    else:
-        return codegen.FunctionCodegen(
-            use_math_symbols=merged_config.use_math_symbols,
-            use_signature=merged_config.use_signature,
-            use_set_symbols=merged_config.use_set_symbols,
-        ).visit(tree)
-
-
-@overload
-def algorithmic(alg: Callable[..., Any], **kwargs: Any) -> output.LatexifiedAlgorithm:
-    ...
-
-
-@overload
-def algorithmic(
-    **kwargs: Any,
-) -> Callable[[Callable[..., Any]], output.LatexifiedAlgorithm]:
-    ...
-
-
-def algorithmic(
-    alg: Callable[..., Any] | None = None, **kwargs: Any
-) -> output.LatexifiedAlgorithm | Callable[
-    [Callable[..., Any]], output.LatexifiedAlgorithm
-]:
-    """Attach LaTeX pretty-printing to the given algorithm.
-
-    This function works with or without specifying the target algorithm as the
-    positional argument. The following two syntaxes works similarly.
-        - latexify.algorithmic(alg, **kwargs)
-        - latexify.algorithmic(**kwargs)(alg)
-
-    Args:
-        alg: Callable to be wrapped.
-        **kwargs: Arguments to control behavior. See also get_latex().
-
-    Returns:
-        - If `alg` is passed, returns the wrapped function.
-        - Otherwise, returns the wrapper function with given settings.
-    """
-    if "style" not in kwargs:
-        kwargs["style"] = Style.ALGORITHMIC
-
-    if alg is not None:
-        return output.LatexifiedAlgorithm(alg, **kwargs)
-
-    def wrapper(a):
-        return output.LatexifiedAlgorithm(a, **kwargs)
-
-    return wrapper
-
-
-@overload
-def function(fn: Callable[..., Any], **kwargs: Any) -> output.LatexifiedFunction:
-    ...
-
-
-@overload
-def function(
-    **kwargs: Any,
-) -> Callable[[Callable[..., Any]], output.LatexifiedFunction]:
-    ...
-
-
-def function(
-    fn: Callable[..., Any] | None = None, **kwargs: Any
-) -> output.LatexifiedFunction | Callable[
-    [Callable[..., Any]], output.LatexifiedFunction
-]:
-    """Attach LaTeX pretty-printing to the given function.
-
-    This function works with or without specifying the target function as the positional
-    argument. The following two syntaxes works similarly.
-        - latexify.function(fn, **kwargs)
-        - latexify.function(**kwargs)(fn)
-
-    Args:
-        fn: Callable to be wrapped.
-        **kwargs: Arguments to control behavior. See also get_latex().
-
-    Returns:
-        - If `fn` is passed, returns the wrapped function.
-        - Otherwise, returns the wrapper function with given settings.
-    """
-    if fn is not None:
-        return output.LatexifiedFunction(fn, **kwargs)
-
-    def wrapper(f):
-        return output.LatexifiedFunction(f, **kwargs)
-
-    return wrapper
-
-
-@overload
-def expression(fn: Callable[..., Any], **kwargs: Any) -> output.LatexifiedFunction:
-    ...
-
-
-@overload
-def expression(
-    **kwargs: Any,
-) -> Callable[[Callable[..., Any]], output.LatexifiedFunction]:
-    ...
-
-
-def expression(
-    fn: Callable[..., Any] | None = None, **kwargs: Any
-) -> output.LatexifiedFunction | Callable[
-    [Callable[..., Any]], output.LatexifiedFunction
-]:
-    """Attach LaTeX pretty-printing to the given function.
-
-    This function is a shortcut for `latexify.function` with the default parameter
-    `use_signature=False`.
-    """
-    kwargs["style"] = Style.EXPRESSION
-    if fn is not None:
-        return function(fn, **kwargs)
-    else:
-        return function(**kwargs)
+import abc
+from typing import Any, Callable
+
+from latexify import exceptions, frontend
+
+
+class LatexifiedRepr(abc.ABC):
+    """Object with LaTeX representation."""
+
+    _fn: Callable[..., Any]
+    _latex: str | None
+    _error: str | None
+
+    def __init__(self, fn, **kwargs):
+        self._fn = fn
+
+    @property
+    def __doc__(self):
+        return self._fn.__doc__
+
+    @__doc__.setter
+    def __doc__(self, val):
+        self._fn.__doc__ = val
+
+    @property
+    def __name__(self):
+        return self._fn.__name__
+
+    @__name__.setter
+    def __name__(self, val):
+        self._fn.__name__ = val
+
+    def __call__(self, *args):
+        return self._fn(*args)
+
+    def __str__(self):
+        return self._latex if self._latex is not None else self._error
+
+    @abc.abstractmethod
+    def _repr_html_(self):
+        """IPython hook to display HTML visualization."""
+        ...
+
+    @abc.abstractmethod
+    def _repr_latex_(self):
+        """IPython hook to display LaTeX visualization."""
+        ...
+
+
+class LatexifiedAlgorithm(LatexifiedRepr):
+    """Algorithm with latex representation."""
+
+    _latex: str | None
+    _error: str | None
+    _jupyter_latex: str | None
+    _jupyter_error: str | None
+
+    def __init__(self, fn, **kwargs):
+        super().__init__(fn)
+        try:
+            kwargs["environment"] = frontend.Environment.LATEX
+            self._latex = frontend.get_latex(fn, **kwargs)
+            self._error = None
+        except exceptions.LatexifyError as e:
+            self._latex = None
+            self._error = f"{type(e).__name__}: {str(e)}"
+        try:
+            kwargs["environment"] = frontend.Environment.JUPYTER_NOTEBOOK
+            self._jupyter_latex = frontend.get_latex(fn, **kwargs)
+            self._jupyter_error = None
+        except exceptions.LatexifyError as e:
+            self._jupyter_latex = None
+            self._jupyter_error = f"{type(e).__name__}: {str(e)}"
+
+    def _repr_html_(self):
+        """IPython hook to display HTML visualization."""
+        return (
+            '<span style="color: red;">' + self._jupyter_error + "</span>"
+            if self._jupyter_error is not None
+            else None
+        )
+
+    def _repr_latex_(self):
+        """IPython hook to display LaTeX visualization."""
+        return (
+            r"$ " + self._jupyter_latex + " $"
+            if self._jupyter_latex is not None
+            else self._jupyter_error
+        )
+
+
+class LatexifiedFunction(LatexifiedRepr):
+    """Function with latex representation."""
+
+    def __init__(self, fn, **kwargs):
+        super().__init__(fn, **kwargs)
+        try:
+            self._latex = frontend.get_latex(fn, **kwargs)
+            self._error = None
+        except exceptions.LatexifyError as e:
+            self._latex = None
+            self._error = f"{type(e).__name__}: {str(e)}"
+
+    def _repr_html_(self):
+        """IPython hook to display HTML visualization."""
+        return (
+            '<span style="color: red;">' + self._error + "</span>"
+            if self._error is not None
+            else None
+        )
+
+    def _repr_latex_(self):
+        """IPython hook to display LaTeX visualization."""
+        return (
+            r"$$ \displaystyle " + self._latex + " $$"
+            if self._latex is not None
+            else self._error
+        )

From fcb47fb1233718fa76cc5a5b0474f8bc971f0425 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 11 Dec 2022 12:07:00 +0000
Subject: [PATCH 11/27] rm dup

---
 src/latexify/output.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/latexify/output.py b/src/latexify/output.py
index 662cdb2..951aeeb 100644
--- a/src/latexify/output.py
+++ b/src/latexify/output.py
@@ -54,8 +54,6 @@ def _repr_latex_(self):
 class LatexifiedAlgorithm(LatexifiedRepr):
     """Algorithm with latex representation."""
 
-    _latex: str | None
-    _error: str | None
     _jupyter_latex: str | None
     _jupyter_error: str | None
 

From ba36b5b3a0fcdea332dbe747ffa83ba03e83232f Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Tue, 13 Dec 2022 20:33:36 +0000
Subject: [PATCH 12/27] init alg

---
 .../algorithmic_style_test.py                 |  50 +++++++
 src/latexify/codegen/__init__.py              |   1 +
 src/latexify/codegen/algorithmic_codegen.py   | 119 +++++++++++++++-
 .../codegen/algorithmic_codegen_test.py       | 128 ++++++++++++++++++
 4 files changed, 297 insertions(+), 1 deletion(-)

diff --git a/src/integration_tests/algorithmic_style_test.py b/src/integration_tests/algorithmic_style_test.py
index 047d6f9..a6d6d7a 100644
--- a/src/integration_tests/algorithmic_style_test.py
+++ b/src/integration_tests/algorithmic_style_test.py
@@ -75,3 +75,53 @@ def collatz(n):
         r" \end{algorithmic}"
     )
     check_algorithm(collatz, latex)
+
+
+def test_factorial_jupyter() -> None:
+    def fact(n):
+        if n == 0:
+            return 1
+        else:
+            return n * fact(n - 1)
+
+    latex = (
+        r"\displaystyle \hspace{0pt} \mathbf{function} \ \mathrm{FACT}(n) \\"
+        r" \displaystyle \hspace{16pt} \mathbf{if} \ n = 0 \\"
+        r" \displaystyle \hspace{32pt} \mathbf{return} \ 1 \\"
+        r" \displaystyle \hspace{16pt} \mathbf{else} \\"
+        r" \displaystyle \hspace{32pt}"
+        r" \mathbf{return} \ n"
+        r" \mathrm{fact} \mathopen{}\left( n - 1 \mathclose{}\right) \\"
+        r" \displaystyle \hspace{16pt} \mathbf{end \ if} \\"
+        r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+    )
+    check_algorithm(fact, latex, environment=frontend.Environment.JUPYTER_NOTEBOOK)
+
+
+def test_collatz_jupyter() -> None:
+    def collatz(n):
+        iterations = 0
+        while n > 1:
+            if n % 2 == 0:
+                n = n // 2
+            else:
+                n = 3 * n + 1
+            iterations = iterations + 1
+        return iterations
+
+    latex = (
+        r"\displaystyle \hspace{0pt} \mathbf{function} \ \mathrm{COLLATZ}(n) \\"
+        r" \displaystyle \hspace{16pt} \mathrm{iterations} \gets 0 \\"
+        r" \displaystyle \hspace{16pt} \mathbf{while} \ n > 1 \\"
+        r" \displaystyle \hspace{32pt} \mathbf{if} \ n \mathbin{\%} 2 = 0 \\"
+        r" \displaystyle \hspace{48pt} n \gets \left\lfloor\frac{n}{2}\right\rfloor \\"
+        r" \displaystyle \hspace{32pt} \mathbf{else} \\"
+        r" \displaystyle \hspace{48pt} n \gets 3 n + 1 \\"
+        r" \displaystyle \hspace{32pt} \mathbf{end \ if} \\"
+        r" \displaystyle \hspace{32pt}"
+        r" \mathrm{iterations} \gets \mathrm{iterations} + 1 \\"
+        r" \displaystyle \hspace{16pt} \mathbf{end \ while} \\"
+        r" \displaystyle \hspace{16pt} \mathbf{return} \ \mathrm{iterations} \\"
+        r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+    )
+    check_algorithm(collatz, latex, environment=frontend.Environment.JUPYTER_NOTEBOOK)
diff --git a/src/latexify/codegen/__init__.py b/src/latexify/codegen/__init__.py
index 1aea2c3..f028adf 100644
--- a/src/latexify/codegen/__init__.py
+++ b/src/latexify/codegen/__init__.py
@@ -3,5 +3,6 @@
 from latexify.codegen import algorithmic_codegen, expression_codegen, function_codegen
 
 AlgorithmicCodegen = algorithmic_codegen.AlgorithmicCodegen
+AlgorithmicJupyterCodegen = algorithmic_codegen.AlgorithmicJupyterCodegen
 ExpressionCodegen = expression_codegen.ExpressionCodegen
 FunctionCodegen = function_codegen.FunctionCodegen
diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index 2fcb16b..973e9df 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -60,6 +60,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         ]
         # Body
         body_strs: list[str] = [self.visit(stmt) for stmt in node.body]
+
         return (
             rf"\begin{{algorithmic}}"
             rf" \Function{{{node.name}}}{{${', '.join(arg_strs)}$}}"
@@ -73,7 +74,6 @@ 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)
-
         latex = rf"\If{{${cond_latex}$}} {body_latex}"
 
         if node.orelse:
@@ -104,3 +104,120 @@ 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"
+
+
+class AlgorithmicJupyterCodegen(ast.NodeVisitor):
+    """Codegen for single algorithms targeting the Jupyter Notebook environment.
+
+    This codegen works for Module with single FunctionDef node to generate a single
+    LaTeX expression of the given algorithm.
+    """
+
+    _PT_PER_INDENT = 16
+    _identifier_converter: identifier_converter.IdentifierConverter
+    _indent: int
+
+    def __init__(
+        self, *, use_math_symbols: bool = False, use_set_symbols: bool = False
+    ) -> None:
+        """Initializer.
+
+        Args:
+            use_math_symbols: Whether to convert identifiers with a math symbol surface
+                (e.g., "alpha") to the LaTeX symbol (e.g., "\\alpha").
+            use_set_symbols: Whether to use set symbols or not.
+        """
+        self._expression_codegen = expression_codegen.ExpressionCodegen(
+            use_math_symbols=use_math_symbols, use_set_symbols=use_set_symbols
+        )
+        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(
+            f"Unsupported AST: {type(node).__name__}"
+        )
+
+    def visit_Assign(self, node: ast.Assign) -> str:
+        """Visit an Assign node."""
+        operands: list[str] = [
+            self._expression_codegen.visit(target) for target in node.targets
+        ]
+        operands.append(self._expression_codegen.visit(node.value))
+        operands_latex = r" \gets ".join(operands)
+        return rf"{self._prefix()} {operands_latex}"
+
+    def visit_Expr(self, node: ast.Expr) -> str:
+        """Visit an Expr node."""
+        return rf"{self._prefix()} {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
+        self._indent += 1
+        body_strs: list[str] = [self.visit(stmt) for stmt in node.body]
+        self._indent -= 1
+        body = r" \\ ".join(body_strs)
+
+        return (
+            rf"{self._prefix()} \mathbf{{function}}"
+            rf" \ \mathrm{{{node.name.upper()}}}({', '.join(arg_strs)}) \\"
+            rf" {body} \\"
+            rf" {self._prefix()} \mathbf{{end \ function}}"
+        )
+
+    # 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
+        body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
+        self._indent -= 1
+        latex = rf"{self._prefix()} \mathbf{{if}} \ {cond_latex} \\ {body_latex}"
+
+        if node.orelse:
+            latex += rf" \\ {self._prefix()} \mathbf{{else}} \\ "
+            self._indent += 1
+            latex += r" \\ ".join(self.visit(stmt) for stmt in node.orelse)
+            self._indent -= 1
+
+        return latex + rf" \\ {self._prefix()} \mathbf{{end \ if}}"
+
+    def visit_Module(self, node: ast.Module) -> str:
+        """Visit a Module node."""
+        return self.visit(node.body[0])
+
+    def visit_Return(self, node: ast.Return) -> str:
+        """Visit a Return node."""
+        return (
+            rf"{self._prefix()} \mathbf{{return}}"
+            rf" \ {self._expression_codegen.visit(node.value)}"
+            if node.value is not None
+            else rf"{self._prefix()} \mathbf{{return}}"
+        )
+
+    def visit_While(self, node: ast.While) -> str:
+        """Visit a While node."""
+        if node.orelse:
+            raise exceptions.LatexifyNotSupportedError(
+                "While statement with the else clause is not supported"
+            )
+
+        cond_latex = self._expression_codegen.visit(node.test)
+        self._indent += 1
+        body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
+        self._indent -= 1
+        return (
+            rf"{self._prefix()} \mathbf{{while}} \ {cond_latex} \\ "
+            rf"{body_latex} \\ "
+            rf"{self._prefix()} \mathbf{{end \ while}}"
+        )
+
+    def _prefix(self) -> str:
+        return rf"\displaystyle \hspace{{{self._indent * self._PT_PER_INDENT}pt}}"
diff --git a/src/latexify/codegen/algorithmic_codegen_test.py b/src/latexify/codegen/algorithmic_codegen_test.py
index a972beb..5051daf 100644
--- a/src/latexify/codegen/algorithmic_codegen_test.py
+++ b/src/latexify/codegen/algorithmic_codegen_test.py
@@ -133,3 +133,131 @@ def test_visit_while_with_else() -> None:
         match="^While statement with the else clause is not supported$",
     ):
         algorithmic_codegen.AlgorithmicCodegen().visit(node)
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        ("x = 3", r"\displaystyle \hspace{0pt} x \gets 3"),
+        ("a = b = 0", r"\displaystyle \hspace{0pt} a \gets b \gets 0"),
+    ],
+)
+def test_visit_assign_jupyter(code: str, latex: str) -> None:
+    node = ast.parse(textwrap.dedent(code)).body[0]
+    assert isinstance(node, ast.Assign)
+    assert algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        (
+            "def f(x): return x",
+            (
+                r"\displaystyle \hspace{0pt} \mathbf{function}"
+                r" \ \mathrm{F}(x) \\"
+                r" \displaystyle \hspace{16pt} \mathbf{return} \ x \\"
+                r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+            ),
+        ),
+        (
+            "def f(a, b, c): return 3",
+            (
+                r"\displaystyle \hspace{0pt} \mathbf{function}"
+                r" \ \mathrm{F}(a, b, c) \\"
+                r" \displaystyle \hspace{16pt} \mathbf{return} \ 3 \\"
+                r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+            ),
+        ),
+    ],
+)
+def test_visit_functiondef_jupyter(code: str, latex: str) -> None:
+    node = ast.parse(textwrap.dedent(code)).body[0]
+    assert isinstance(node, ast.FunctionDef)
+    assert algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        (
+            "if x < y: return x",
+            (
+                r"\displaystyle \hspace{0pt} \mathbf{if} \ x < y \\"
+                r" \displaystyle \hspace{16pt} \mathbf{return} \ x \\"
+                r" \displaystyle \hspace{0pt} \mathbf{end \ if}"
+            ),
+        ),
+        (
+            "if True: x\nelse: y",
+            (
+                r"\displaystyle \hspace{0pt} \mathbf{if} \ \mathrm{True} \\"
+                r" \displaystyle \hspace{16pt} x \\"
+                r" \displaystyle \hspace{0pt} \mathbf{else} \\"
+                r" \displaystyle \hspace{16pt} y \\"
+                r" \displaystyle \hspace{0pt} \mathbf{end \ if}"
+            ),
+        ),
+    ],
+)
+def test_visit_if_jupyter(code: str, latex: str) -> None:
+    node = ast.parse(textwrap.dedent(code)).body[0]
+    assert isinstance(node, ast.If)
+    assert algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        (
+            "return x + y",
+            r"\displaystyle \hspace{0pt} \mathbf{return} \ x + y",
+        ),
+        (
+            "return",
+            r"\displaystyle \hspace{0pt} \mathbf{return}",
+        ),
+    ],
+)
+def test_visit_return_jupyter(code: str, latex: str) -> None:
+    node = ast.parse(textwrap.dedent(code)).body[0]
+    assert isinstance(node, ast.Return)
+    assert algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        (
+            "while x < y: x = x + 1",
+            (
+                r"\displaystyle \hspace{0pt} \mathbf{while} \ x < y \\"
+                r" \displaystyle \hspace{16pt} x \gets x + 1 \\"
+                r" \displaystyle \hspace{0pt} \mathbf{end \ while}"
+            ),
+        )
+    ],
+)
+def test_visit_while_jupyter(code: str, latex: str) -> None:
+    node = ast.parse(textwrap.dedent(code)).body[0]
+    assert isinstance(node, ast.While)
+    assert algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node) == latex
+
+
+def test_visit_while_with_else_jupyter() -> None:
+    node = ast.parse(
+        textwrap.dedent(
+            """
+            while True:
+                x = x
+            else:
+                x = y
+            """
+        )
+    ).body[0]
+    assert isinstance(node, ast.While)
+    with pytest.raises(
+        exceptions.LatexifyNotSupportedError,
+        match="^While statement with the else clause is not supported$",
+    ):
+        algorithmic_codegen.AlgorithmicJupyterCodegen().visit(node)

From 006821df10e61ee2c5d22b90fb7fd14916d954f5 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Tue, 13 Dec 2022 21:49:05 +0000
Subject: [PATCH 13/27] wrap up

---
 .../algorithmic_style_test.py                 | 39 +++++++-------
 src/latexify/codegen/algorithmic_codegen.py   | 30 ++++++-----
 .../codegen/algorithmic_codegen_test.py       | 52 +++++++++----------
 3 files changed, 63 insertions(+), 58 deletions(-)

diff --git a/src/integration_tests/algorithmic_style_test.py b/src/integration_tests/algorithmic_style_test.py
index 4b7686b..4f6121b 100644
--- a/src/integration_tests/algorithmic_style_test.py
+++ b/src/integration_tests/algorithmic_style_test.py
@@ -93,15 +93,15 @@ def fact(n):
             return n * fact(n - 1)
 
     latex = (
-        r"\displaystyle \hspace{0pt} \mathbf{function} \ \mathrm{FACT}(n) \\"
-        r" \displaystyle \hspace{16pt} \mathbf{if} \ n = 0 \\"
-        r" \displaystyle \hspace{32pt} \mathbf{return} \ 1 \\"
-        r" \displaystyle \hspace{16pt} \mathbf{else} \\"
-        r" \displaystyle \hspace{32pt}"
+        r"\mathbf{function} \ \mathrm{FACT}(n) \\"
+        r" \hspace{1em} \mathbf{if} \ n = 0 \\"
+        r" \hspace{2em} \mathbf{return} \ 1 \\"
+        r" \hspace{1em} \mathbf{else} \\"
+        r" \hspace{2em}"
         r" \mathbf{return} \ n \cdot"
         r" \mathrm{fact} \mathopen{}\left( n - 1 \mathclose{}\right) \\"
-        r" \displaystyle \hspace{16pt} \mathbf{end \ if} \\"
-        r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+        r" \hspace{1em} \mathbf{end \ if} \\"
+        r" \mathbf{end \ function}"
     )
     check_algorithm(fact, latex, generate_latex.Style.IPYTHON_ALGORITHMIC)
 
@@ -118,18 +118,19 @@ def collatz(n):
         return iterations
 
     latex = (
-        r"\displaystyle \hspace{0pt} \mathbf{function} \ \mathrm{COLLATZ}(n) \\"
-        r" \displaystyle \hspace{16pt} \mathrm{iterations} \gets 0 \\"
-        r" \displaystyle \hspace{16pt} \mathbf{while} \ n > 1 \\"
-        r" \displaystyle \hspace{32pt} \mathbf{if} \ n \mathbin{\%} 2 = 0 \\"
-        r" \displaystyle \hspace{48pt} n \gets \left\lfloor\frac{n}{2}\right\rfloor \\"
-        r" \displaystyle \hspace{32pt} \mathbf{else} \\"
-        r" \displaystyle \hspace{48pt} n \gets 3 \cdot n + 1 \\"
-        r" \displaystyle \hspace{32pt} \mathbf{end \ if} \\"
-        r" \displaystyle \hspace{32pt}"
+        r"\mathbf{function} \ \mathrm{COLLATZ}(n) \\"
+        r" \hspace{1em} \mathrm{iterations} \gets 0 \\"
+        r" \hspace{1em} \mathbf{while} \ n > 1 \\"
+        r" \hspace{2em} \mathbf{if} \ n \mathbin{\%} 2 = 0 \\"
+        r" \hspace{3em} n \gets \left\lfloor\frac{n}{2}\right\rfloor \\"
+        r" \hspace{2em} \mathbf{else} \\"
+        r" \hspace{3em} n \gets 3 \cdot n + 1 \\"
+        r" \hspace{2em} \mathbf{end \ if} \\"
+        r" \hspace{2em}"
         r" \mathrm{iterations} \gets \mathrm{iterations} + 1 \\"
-        r" \displaystyle \hspace{16pt} \mathbf{end \ while} \\"
-        r" \displaystyle \hspace{16pt} \mathbf{return} \ \mathrm{iterations} \\"
-        r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+        r" \hspace{1em} \mathbf{end \ while} \\"
+        r" \hspace{1em} \mathbf{return} \ \mathrm{iterations} \\"
+        r" \mathbf{end \ function}"
     )
+
     check_algorithm(collatz, latex, generate_latex.Style.IPYTHON_ALGORITHMIC)
diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index 3a4ed7f..0e09db7 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -151,7 +151,7 @@ class IPythonAlgorithmicCodegen(ast.NodeVisitor):
     LaTeX expression of the given algorithm.
     """
 
-    _PT_PER_INDENT = 16
+    _EM_PER_INDENT = 1
 
     _identifier_converter: identifier_converter.IdentifierConverter
     _indent_level: int
@@ -186,11 +186,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"{self._add_prefix()} {operands_latex}"
+        return rf"{self._add_prefix()}{operands_latex}"
 
     def visit_Expr(self, node: ast.Expr) -> str:
         """Visit an Expr node."""
-        return rf"{self._add_prefix()} {self._expression_codegen.visit(node.value)}"
+        return rf"{self._add_prefix()}{self._expression_codegen.visit(node.value)}"
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         """Visit a FunctionDef node."""
@@ -204,10 +204,10 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         body = r" \\ ".join(body_strs)
 
         return (
-            rf"{self._add_prefix()} \mathbf{{function}}"
+            rf"{self._add_prefix()}\mathbf{{function}}"
             rf" \ \mathrm{{{node.name.upper()}}}({', '.join(arg_strs)}) \\"
             rf" {body} \\"
-            rf" {self._add_prefix()} \mathbf{{end \ function}}"
+            rf" {self._add_prefix()}\mathbf{{end \ function}}"
         )
 
     # TODO(ZibingZhang): support \ELSIF
@@ -216,14 +216,14 @@ def visit_If(self, node: ast.If) -> str:
         cond_latex = self._expression_codegen.visit(node.test)
         with self._increment_level():
             body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
-        latex = rf"{self._add_prefix()} \mathbf{{if}} \ {cond_latex} \\ {body_latex}"
+        latex = rf"{self._add_prefix()}\mathbf{{if}} \ {cond_latex} \\ {body_latex}"
 
         if node.orelse:
-            latex += rf" \\ {self._add_prefix()} \mathbf{{else}} \\ "
+            latex += rf" \\ {self._add_prefix()}\mathbf{{else}} \\ "
             with self._increment_level():
                 latex += r" \\ ".join(self.visit(stmt) for stmt in node.orelse)
 
-        return latex + rf" \\ {self._add_prefix()} \mathbf{{end \ if}}"
+        return latex + rf" \\ {self._add_prefix()}\mathbf{{end \ if}}"
 
     def visit_Module(self, node: ast.Module) -> str:
         """Visit a Module node."""
@@ -232,10 +232,10 @@ def visit_Module(self, node: ast.Module) -> str:
     def visit_Return(self, node: ast.Return) -> str:
         """Visit a Return node."""
         return (
-            rf"{self._add_prefix()} \mathbf{{return}}"
+            rf"{self._add_prefix()}\mathbf{{return}}"
             rf" \ {self._expression_codegen.visit(node.value)}"
             if node.value is not None
-            else rf"{self._add_prefix()} \mathbf{{return}}"
+            else rf"{self._add_prefix()}\mathbf{{return}}"
         )
 
     def visit_While(self, node: ast.While) -> str:
@@ -249,9 +249,9 @@ def visit_While(self, node: ast.While) -> str:
         with self._increment_level():
             body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
         return (
-            rf"{self._add_prefix()} \mathbf{{while}} \ {cond_latex} \\ "
+            rf"{self._add_prefix()}\mathbf{{while}} \ {cond_latex} \\ "
             rf"{body_latex} \\ "
-            rf"{self._add_prefix()} \mathbf{{end \ while}}"
+            rf"{self._add_prefix()}\mathbf{{end \ while}}"
         )
 
     @contextlib.contextmanager
@@ -262,4 +262,8 @@ def _increment_level(self) -> Generator[None, None, None]:
         self._indent_level -= 1
 
     def _add_prefix(self) -> str:
-        return rf"\displaystyle \hspace{{{self._indent_level * self._PT_PER_INDENT}pt}}"
+        return (
+            rf"\hspace{{{self._indent_level * self._EM_PER_INDENT}em}} "
+            if self._indent_level > 0
+            else ""
+        )
diff --git a/src/latexify/codegen/algorithmic_codegen_test.py b/src/latexify/codegen/algorithmic_codegen_test.py
index f04890b..6fd8a59 100644
--- a/src/latexify/codegen/algorithmic_codegen_test.py
+++ b/src/latexify/codegen/algorithmic_codegen_test.py
@@ -170,8 +170,8 @@ def test_visit_while_with_else() -> None:
 @pytest.mark.parametrize(
     "code,latex",
     [
-        ("x = 3", r"\displaystyle \hspace{0pt} x \gets 3"),
-        ("a = b = 0", r"\displaystyle \hspace{0pt} a \gets b \gets 0"),
+        ("x = 3", r"x \gets 3"),
+        ("a = b = 0", r"a \gets b \gets 0"),
     ],
 )
 def test_visit_assign_jupyter(code: str, latex: str) -> None:
@@ -186,24 +186,24 @@ def test_visit_assign_jupyter(code: str, latex: str) -> None:
         (
             "def f(x): return x",
             (
-                r"\displaystyle \hspace{0pt} \mathbf{function}"
+                r"\mathbf{function}"
                 r" \ \mathrm{F}(x) \\"
-                r" \displaystyle \hspace{16pt} \mathbf{return} \ x \\"
-                r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+                r" \hspace{1em} \mathbf{return} \ x \\"
+                r" \mathbf{end \ function}"
             ),
         ),
         (
             "def f(a, b, c): return 3",
             (
-                r"\displaystyle \hspace{0pt} \mathbf{function}"
+                r"\mathbf{function}"
                 r" \ \mathrm{F}(a, b, c) \\"
-                r" \displaystyle \hspace{16pt} \mathbf{return} \ 3 \\"
-                r" \displaystyle \hspace{0pt} \mathbf{end \ function}"
+                r" \hspace{1em} \mathbf{return} \ 3 \\"
+                r" \mathbf{end \ function}"
             ),
         ),
     ],
 )
-def test_visit_functiondef_jupyter(code: str, latex: str) -> None:
+def test_visit_functiondef_ipython(code: str, latex: str) -> None:
     node = ast.parse(textwrap.dedent(code)).body[0]
     assert isinstance(node, ast.FunctionDef)
     assert algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == latex
@@ -215,24 +215,24 @@ def test_visit_functiondef_jupyter(code: str, latex: str) -> None:
         (
             "if x < y: return x",
             (
-                r"\displaystyle \hspace{0pt} \mathbf{if} \ x < y \\"
-                r" \displaystyle \hspace{16pt} \mathbf{return} \ x \\"
-                r" \displaystyle \hspace{0pt} \mathbf{end \ if}"
+                r"\mathbf{if} \ x < y \\"
+                r" \hspace{1em} \mathbf{return} \ x \\"
+                r" \mathbf{end \ if}"
             ),
         ),
         (
             "if True: x\nelse: y",
             (
-                r"\displaystyle \hspace{0pt} \mathbf{if} \ \mathrm{True} \\"
-                r" \displaystyle \hspace{16pt} x \\"
-                r" \displaystyle \hspace{0pt} \mathbf{else} \\"
-                r" \displaystyle \hspace{16pt} y \\"
-                r" \displaystyle \hspace{0pt} \mathbf{end \ if}"
+                r"\mathbf{if} \ \mathrm{True} \\"
+                r" \hspace{1em} x \\"
+                r" \mathbf{else} \\"
+                r" \hspace{1em} y \\"
+                r" \mathbf{end \ if}"
             ),
         ),
     ],
 )
-def test_visit_if_jupyter(code: str, latex: str) -> None:
+def test_visit_if_ipython(code: str, latex: str) -> None:
     node = ast.parse(textwrap.dedent(code)).body[0]
     assert isinstance(node, ast.If)
     assert algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == latex
@@ -243,15 +243,15 @@ def test_visit_if_jupyter(code: str, latex: str) -> None:
     [
         (
             "return x + y",
-            r"\displaystyle \hspace{0pt} \mathbf{return} \ x + y",
+            r"\mathbf{return} \ x + y",
         ),
         (
             "return",
-            r"\displaystyle \hspace{0pt} \mathbf{return}",
+            r"\mathbf{return}",
         ),
     ],
 )
-def test_visit_return_jupyter(code: str, latex: str) -> None:
+def test_visit_return_ipython(code: str, latex: str) -> None:
     node = ast.parse(textwrap.dedent(code)).body[0]
     assert isinstance(node, ast.Return)
     assert algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == latex
@@ -263,20 +263,20 @@ def test_visit_return_jupyter(code: str, latex: str) -> None:
         (
             "while x < y: x = x + 1",
             (
-                r"\displaystyle \hspace{0pt} \mathbf{while} \ x < y \\"
-                r" \displaystyle \hspace{16pt} x \gets x + 1 \\"
-                r" \displaystyle \hspace{0pt} \mathbf{end \ while}"
+                r"\mathbf{while} \ x < y \\"
+                r" \hspace{1em} x \gets x + 1 \\"
+                r" \mathbf{end \ while}"
             ),
         )
     ],
 )
-def test_visit_while_jupyter(code: str, latex: str) -> None:
+def test_visit_while_ipython(code: str, latex: str) -> None:
     node = ast.parse(textwrap.dedent(code)).body[0]
     assert isinstance(node, ast.While)
     assert algorithmic_codegen.IPythonAlgorithmicCodegen().visit(node) == latex
 
 
-def test_visit_while_with_else_jupyter() -> None:
+def test_visit_while_with_else_ipython() -> None:
     node = ast.parse(
         textwrap.dedent(
             """

From dafa24f77b2e4ea6288288e3ecf6ddd9af4c3f4e Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Tue, 13 Dec 2022 21:50:51 +0000
Subject: [PATCH 14/27] oop

---
 src/latexify/codegen/__init__.py |   2 +-
 src/latexify/output.py           | 120 -------------------------------
 2 files changed, 1 insertion(+), 121 deletions(-)
 delete mode 100644 src/latexify/output.py

diff --git a/src/latexify/codegen/__init__.py b/src/latexify/codegen/__init__.py
index 1b0940a..8d09290 100644
--- a/src/latexify/codegen/__init__.py
+++ b/src/latexify/codegen/__init__.py
@@ -3,6 +3,6 @@
 from latexify.codegen import algorithmic_codegen, expression_codegen, function_codegen
 
 AlgorithmicCodegen = algorithmic_codegen.AlgorithmicCodegen
-IPythonAlgorithmicCodegen = algorithmic_codegen.IPythonAlgorithmicCodegen
 ExpressionCodegen = expression_codegen.ExpressionCodegen
 FunctionCodegen = function_codegen.FunctionCodegen
+IPythonAlgorithmicCodegen = algorithmic_codegen.IPythonAlgorithmicCodegen
diff --git a/src/latexify/output.py b/src/latexify/output.py
deleted file mode 100644
index 951aeeb..0000000
--- a/src/latexify/output.py
+++ /dev/null
@@ -1,120 +0,0 @@
-"""Output of the frontend function decorators."""
-
-from __future__ import annotations
-
-import abc
-from typing import Any, Callable
-
-from latexify import exceptions, frontend
-
-
-class LatexifiedRepr(abc.ABC):
-    """Object with LaTeX representation."""
-
-    _fn: Callable[..., Any]
-    _latex: str | None
-    _error: str | None
-
-    def __init__(self, fn, **kwargs):
-        self._fn = fn
-
-    @property
-    def __doc__(self):
-        return self._fn.__doc__
-
-    @__doc__.setter
-    def __doc__(self, val):
-        self._fn.__doc__ = val
-
-    @property
-    def __name__(self):
-        return self._fn.__name__
-
-    @__name__.setter
-    def __name__(self, val):
-        self._fn.__name__ = val
-
-    def __call__(self, *args):
-        return self._fn(*args)
-
-    def __str__(self):
-        return self._latex if self._latex is not None else self._error
-
-    @abc.abstractmethod
-    def _repr_html_(self):
-        """IPython hook to display HTML visualization."""
-        ...
-
-    @abc.abstractmethod
-    def _repr_latex_(self):
-        """IPython hook to display LaTeX visualization."""
-        ...
-
-
-class LatexifiedAlgorithm(LatexifiedRepr):
-    """Algorithm with latex representation."""
-
-    _jupyter_latex: str | None
-    _jupyter_error: str | None
-
-    def __init__(self, fn, **kwargs):
-        super().__init__(fn)
-        try:
-            kwargs["environment"] = frontend.Environment.LATEX
-            self._latex = frontend.get_latex(fn, **kwargs)
-            self._error = None
-        except exceptions.LatexifyError as e:
-            self._latex = None
-            self._error = f"{type(e).__name__}: {str(e)}"
-        try:
-            kwargs["environment"] = frontend.Environment.JUPYTER_NOTEBOOK
-            self._jupyter_latex = frontend.get_latex(fn, **kwargs)
-            self._jupyter_error = None
-        except exceptions.LatexifyError as e:
-            self._jupyter_latex = None
-            self._jupyter_error = f"{type(e).__name__}: {str(e)}"
-
-    def _repr_html_(self):
-        """IPython hook to display HTML visualization."""
-        return (
-            '<span style="color: red;">' + self._jupyter_error + "</span>"
-            if self._jupyter_error is not None
-            else None
-        )
-
-    def _repr_latex_(self):
-        """IPython hook to display LaTeX visualization."""
-        return (
-            r"$ " + self._jupyter_latex + " $"
-            if self._jupyter_latex is not None
-            else self._jupyter_error
-        )
-
-
-class LatexifiedFunction(LatexifiedRepr):
-    """Function with latex representation."""
-
-    def __init__(self, fn, **kwargs):
-        super().__init__(fn, **kwargs)
-        try:
-            self._latex = frontend.get_latex(fn, **kwargs)
-            self._error = None
-        except exceptions.LatexifyError as e:
-            self._latex = None
-            self._error = f"{type(e).__name__}: {str(e)}"
-
-    def _repr_html_(self):
-        """IPython hook to display HTML visualization."""
-        return (
-            '<span style="color: red;">' + self._error + "</span>"
-            if self._error is not None
-            else None
-        )
-
-    def _repr_latex_(self):
-        """IPython hook to display LaTeX visualization."""
-        return (
-            r"$$ \displaystyle " + self._latex + " $$"
-            if self._latex is not None
-            else self._error
-        )

From ffe0a14e33dc584d6932630e8a844b24959ba415 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Tue, 13 Dec 2022 22:03:40 +0000
Subject: [PATCH 15/27] testing

---
 .../algorithmic_style_test.py                 | 78 ++++---------------
 src/integration_tests/integration_utils.py    | 40 ++++++++++
 src/latexify/codegen/algorithmic_codegen.py   |  4 +-
 src/latexify/ipython_wrappers.py              |  2 +-
 4 files changed, 58 insertions(+), 66 deletions(-)

diff --git a/src/integration_tests/algorithmic_style_test.py b/src/integration_tests/algorithmic_style_test.py
index 4f6121b..9a0cf46 100644
--- a/src/integration_tests/algorithmic_style_test.py
+++ b/src/integration_tests/algorithmic_style_test.py
@@ -3,31 +3,8 @@
 from __future__ import annotations
 
 import textwrap
-from typing import Any, Callable
 
-from latexify import generate_latex
-
-
-def check_algorithm(
-    fn: Callable[..., Any],
-    latex: str,
-    style: generate_latex.Style,
-    **kwargs,
-) -> None:
-    """Helper to check if the obtained function has the expected LaTeX form.
-
-    Args:
-        fn: Function to check.
-        latex: LaTeX form of `fn`.
-        style: The style of the output.
-        **kwargs: Arguments passed to `frontend.get_latex`.
-    """
-    # Checks the syntax:
-    #     def fn(...):
-    #         ...
-    #     latexified = get_latex(fn, style=ALGORITHM, **kwargs)
-    latexified = generate_latex.get_latex(fn, style=style, **kwargs)
-    assert latexified == latex
+from integration_tests import integration_utils
 
 
 def test_factorial() -> None:
@@ -50,7 +27,18 @@ def fact(n):
         \end{algorithmic}
         """  # noqa: E501
     ).strip()
-    check_algorithm(fact, latex, generate_latex.Style.ALGORITHMIC)
+    ipython_latex = (
+        r"\mathbf{function} \ \mathrm{FACT}(n) \\"
+        r" \hspace{1em} \mathbf{if} \ n = 0 \\"
+        r" \hspace{2em} \mathbf{return} \ 1 \\"
+        r" \hspace{1em} \mathbf{else} \\"
+        r" \hspace{2em}"
+        r" \mathbf{return} \ n \cdot"
+        r" \mathrm{fact} \mathopen{}\left( n - 1 \mathclose{}\right) \\"
+        r" \hspace{1em} \mathbf{end \ if} \\"
+        r" \mathbf{end \ function}"
+    )
+    integration_utils.check_algorithm(fact, latex, ipython_latex)
 
 
 def test_collatz() -> None:
@@ -82,42 +70,7 @@ def collatz(n):
         \end{algorithmic}
         """
     ).strip()
-    check_algorithm(collatz, latex, generate_latex.Style.ALGORITHMIC)
-
-
-def test_factorial_jupyter() -> None:
-    def fact(n):
-        if n == 0:
-            return 1
-        else:
-            return n * fact(n - 1)
-
-    latex = (
-        r"\mathbf{function} \ \mathrm{FACT}(n) \\"
-        r" \hspace{1em} \mathbf{if} \ n = 0 \\"
-        r" \hspace{2em} \mathbf{return} \ 1 \\"
-        r" \hspace{1em} \mathbf{else} \\"
-        r" \hspace{2em}"
-        r" \mathbf{return} \ n \cdot"
-        r" \mathrm{fact} \mathopen{}\left( n - 1 \mathclose{}\right) \\"
-        r" \hspace{1em} \mathbf{end \ if} \\"
-        r" \mathbf{end \ function}"
-    )
-    check_algorithm(fact, latex, generate_latex.Style.IPYTHON_ALGORITHMIC)
-
-
-def test_collatz_jupyter() -> None:
-    def collatz(n):
-        iterations = 0
-        while n > 1:
-            if n % 2 == 0:
-                n = n // 2
-            else:
-                n = 3 * n + 1
-            iterations = iterations + 1
-        return iterations
-
-    latex = (
+    ipython_latex = (
         r"\mathbf{function} \ \mathrm{COLLATZ}(n) \\"
         r" \hspace{1em} \mathrm{iterations} \gets 0 \\"
         r" \hspace{1em} \mathbf{while} \ n > 1 \\"
@@ -132,5 +85,4 @@ def collatz(n):
         r" \hspace{1em} \mathbf{return} \ \mathrm{iterations} \\"
         r" \mathbf{end \ function}"
     )
-
-    check_algorithm(collatz, latex, generate_latex.Style.IPYTHON_ALGORITHMIC)
+    integration_utils.check_algorithm(collatz, latex, ipython_latex)
diff --git a/src/integration_tests/integration_utils.py b/src/integration_tests/integration_utils.py
index 92ada9a..5ffe7a0 100644
--- a/src/integration_tests/integration_utils.py
+++ b/src/integration_tests/integration_utils.py
@@ -43,3 +43,43 @@ def check_function(
     latexified = frontend.function(fn, **kwargs)
     assert str(latexified) == latex
     assert latexified._repr_latex_() == rf"$$ \displaystyle {latex} $$"
+
+
+def check_algorithm(
+    fn: Callable[..., Any],
+    latex: str,
+    ipython_latex: str,
+    **kwargs,
+) -> None:
+    """Helper to check if the obtained function has the expected LaTeX form.
+
+    Args:
+        fn: Function to check.
+        latex: LaTeX form of `fn`.
+        ipython_latex: IPython LaTeX form of `fn`
+        **kwargs: Arguments passed to `frontend.get_latex`.
+    """
+    # Checks the syntax:
+    #     @algorithmic
+    #     def fn(...):
+    #         ...
+    if not kwargs:
+        latexified = frontend.algorithmic(fn)
+        assert str(latexified) == latex
+        assert latexified._repr_latex_() == f"$ {ipython_latex} $"
+
+    # Checks the syntax:
+    #     @algorithmic(**kwargs)
+    #     def fn(...):
+    #         ...
+    latexified = frontend.algorithmic(**kwargs)(fn)
+    assert str(latexified) == latex
+    assert latexified._repr_latex_() == f"$ {ipython_latex} $"
+
+    # Checks the syntax:
+    #     def fn(...):
+    #         ...
+    #     latexified = algorithmic(fn, **kwargs)
+    latexified = frontend.algorithmic(fn, **kwargs)
+    assert str(latexified) == latex
+    assert latexified._repr_latex_() == f"$ {ipython_latex} $"
diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index 0e09db7..4faa57c 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -186,11 +186,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"{self._add_prefix()}{operands_latex}"
+        return self._add_prefix() + operands_latex
 
     def visit_Expr(self, node: ast.Expr) -> str:
         """Visit an Expr node."""
-        return rf"{self._add_prefix()}{self._expression_codegen.visit(node.value)}"
+        return self._add_prefix() + self._expression_codegen.visit(node.value)
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         """Visit a FunctionDef node."""
diff --git a/src/latexify/ipython_wrappers.py b/src/latexify/ipython_wrappers.py
index 88aa015..18be29a 100644
--- a/src/latexify/ipython_wrappers.py
+++ b/src/latexify/ipython_wrappers.py
@@ -95,7 +95,7 @@ def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:
     def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:
         """IPython hook to display LaTeX visualization."""
         return (
-            r"$ " + self._ipython_latex + " $"
+            f"$ {self._ipython_latex} $"
             if self._ipython_latex is not None
             else self._ipython_error
         )

From 2a5a0b9d8b905c7a1a4ba9ba80ec1b7ab3cb5ddc Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Tue, 13 Dec 2022 22:05:12 +0000
Subject: [PATCH 16/27] change back

---
 src/latexify/ipython_wrappers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/latexify/ipython_wrappers.py b/src/latexify/ipython_wrappers.py
index 18be29a..ef4c482 100644
--- a/src/latexify/ipython_wrappers.py
+++ b/src/latexify/ipython_wrappers.py
@@ -95,7 +95,7 @@ def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:
     def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:
         """IPython hook to display LaTeX visualization."""
         return (
-            f"$ {self._ipython_latex} $"
+            "$ " + self._ipython_latex + " $"
             if self._ipython_latex is not None
             else self._ipython_error
         )

From 7fbd04896f33c6048e6ffd90ed15f8697b21f9e9 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 09:04:34 +0000
Subject: [PATCH 17/27] suggestions

---
 src/integration_tests/algorithmic_style_test.py  |  8 ++++++--
 src/latexify/codegen/algorithmic_codegen.py      |  8 ++++++--
 src/latexify/codegen/algorithmic_codegen_test.py | 12 ++++++++----
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/integration_tests/algorithmic_style_test.py b/src/integration_tests/algorithmic_style_test.py
index 9a0cf46..dbe2b5d 100644
--- a/src/integration_tests/algorithmic_style_test.py
+++ b/src/integration_tests/algorithmic_style_test.py
@@ -28,7 +28,8 @@ def fact(n):
         """  # noqa: E501
     ).strip()
     ipython_latex = (
-        r"\mathbf{function} \ \mathrm{FACT}(n) \\"
+        r"\begin{array}{l}"
+        r" \mathbf{function} \ \texttt{FACT}(n) \\"
         r" \hspace{1em} \mathbf{if} \ n = 0 \\"
         r" \hspace{2em} \mathbf{return} \ 1 \\"
         r" \hspace{1em} \mathbf{else} \\"
@@ -37,6 +38,7 @@ def fact(n):
         r" \mathrm{fact} \mathopen{}\left( n - 1 \mathclose{}\right) \\"
         r" \hspace{1em} \mathbf{end \ if} \\"
         r" \mathbf{end \ function}"
+        r" \end{array}"
     )
     integration_utils.check_algorithm(fact, latex, ipython_latex)
 
@@ -71,7 +73,8 @@ def collatz(n):
         """
     ).strip()
     ipython_latex = (
-        r"\mathbf{function} \ \mathrm{COLLATZ}(n) \\"
+        r"\begin{array}{l}"
+        r" \mathbf{function} \ \texttt{COLLATZ}(n) \\"
         r" \hspace{1em} \mathrm{iterations} \gets 0 \\"
         r" \hspace{1em} \mathbf{while} \ n > 1 \\"
         r" \hspace{2em} \mathbf{if} \ n \mathbin{\%} 2 = 0 \\"
@@ -84,5 +87,6 @@ def collatz(n):
         r" \hspace{1em} \mathbf{end \ while} \\"
         r" \hspace{1em} \mathbf{return} \ \mathrm{iterations} \\"
         r" \mathbf{end \ function}"
+        r" \end{array}"
     )
     integration_utils.check_algorithm(collatz, latex, ipython_latex)
diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index 4faa57c..18ce407 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -60,6 +60,7 @@ def visit_Expr(self, node: ast.Expr) -> str:
             rf"\State ${self._expression_codegen.visit(node.value)}$"
         )
 
+    # TODO(ZibingZhang): support nested functions
     def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         """Visit a FunctionDef node."""
         # Arguments
@@ -192,6 +193,7 @@ def visit_Expr(self, node: ast.Expr) -> str:
         """Visit an Expr node."""
         return self._add_prefix() + self._expression_codegen.visit(node.value)
 
+    # TODO(ZibingZhang): support nested functions
     def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         """Visit a FunctionDef node."""
         # Arguments
@@ -204,10 +206,12 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         body = r" \\ ".join(body_strs)
 
         return (
-            rf"{self._add_prefix()}\mathbf{{function}}"
-            rf" \ \mathrm{{{node.name.upper()}}}({', '.join(arg_strs)}) \\"
+            r"\begin{array}{l}"
+            rf" {self._add_prefix()}\mathbf{{function}}"
+            rf" \ \texttt{{{node.name.upper()}}}({', '.join(arg_strs)}) \\"
             rf" {body} \\"
             rf" {self._add_prefix()}\mathbf{{end \ function}}"
+            r" \end{array}"
         )
 
     # TODO(ZibingZhang): support \ELSIF
diff --git a/src/latexify/codegen/algorithmic_codegen_test.py b/src/latexify/codegen/algorithmic_codegen_test.py
index 6fd8a59..58d8244 100644
--- a/src/latexify/codegen/algorithmic_codegen_test.py
+++ b/src/latexify/codegen/algorithmic_codegen_test.py
@@ -186,19 +186,23 @@ def test_visit_assign_jupyter(code: str, latex: str) -> None:
         (
             "def f(x): return x",
             (
-                r"\mathbf{function}"
-                r" \ \mathrm{F}(x) \\"
+                r"\begin{array}{l}"
+                r" \mathbf{function}"
+                r" \ \texttt{F}(x) \\"
                 r" \hspace{1em} \mathbf{return} \ x \\"
                 r" \mathbf{end \ function}"
+                r" \end{array}"
             ),
         ),
         (
             "def f(a, b, c): return 3",
             (
-                r"\mathbf{function}"
-                r" \ \mathrm{F}(a, b, c) \\"
+                r"\begin{array}{l}"
+                r" \mathbf{function}"
+                r" \ \texttt{F}(a, b, c) \\"
                 r" \hspace{1em} \mathbf{return} \ 3 \\"
                 r" \mathbf{end \ function}"
+                r" \end{array}"
             ),
         ),
     ],

From 65f8ab80f5d0b694515293a24ec25d885e189046 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 09:10:52 +0000
Subject: [PATCH 18/27] f strings >

---
 src/latexify/ipython_wrappers.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/latexify/ipython_wrappers.py b/src/latexify/ipython_wrappers.py
index ef4c482..c77a869 100644
--- a/src/latexify/ipython_wrappers.py
+++ b/src/latexify/ipython_wrappers.py
@@ -95,7 +95,7 @@ def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:
     def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:
         """IPython hook to display LaTeX visualization."""
         return (
-            "$ " + self._ipython_latex + " $"
+            f"$ {self._ipython_latex} $"
             if self._ipython_latex is not None
             else self._ipython_error
         )
@@ -133,7 +133,7 @@ def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:
     def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:
         """IPython hook to display LaTeX visualization."""
         return (
-            r"$$ \displaystyle " + self._latex + " $$"
+            rf"$$ \displaystyle {self._latex} $$"
             if self._latex is not None
             else self._error
         )

From fb7612a67626f82dc40993b95e4bc4e4d52121ac Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 09:19:25 +0000
Subject: [PATCH 19/27] slight refactor

---
 src/latexify/codegen/algorithmic_codegen.py | 53 ++++++++++++---------
 1 file changed, 31 insertions(+), 22 deletions(-)

diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index 18ce407..a2fdd7b 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -137,10 +137,10 @@ def _increment_level(self) -> Generator[None, None, None]:
         self._indent_level -= 1
 
     def _add_indent(self, line: str) -> str:
-        """Adds whitespace before the line.
+        """Adds an indent before the line.
 
         Args:
-            line: The line to add whitespace to.
+            line: The line to add an indent to.
         """
         return self._indent_level * self._SPACES_PER_INDENT * " " + line
 
@@ -187,11 +187,11 @@ def visit_Assign(self, node: ast.Assign) -> str:
         ]
         operands.append(self._expression_codegen.visit(node.value))
         operands_latex = r" \gets ".join(operands)
-        return self._add_prefix() + operands_latex
+        return self._add_indent(operands_latex)
 
     def visit_Expr(self, node: ast.Expr) -> str:
         """Visit an Expr node."""
-        return self._add_prefix() + self._expression_codegen.visit(node.value)
+        return self._add_indent(self._expression_codegen.visit(node.value))
 
     # TODO(ZibingZhang): support nested functions
     def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
@@ -206,12 +206,13 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         body = r" \\ ".join(body_strs)
 
         return (
-            r"\begin{array}{l}"
-            rf" {self._add_prefix()}\mathbf{{function}}"
-            rf" \ \texttt{{{node.name.upper()}}}({', '.join(arg_strs)}) \\"
-            rf" {body} \\"
-            rf" {self._add_prefix()}\mathbf{{end \ function}}"
-            r" \end{array}"
+            r"\begin{array}{l} "
+            + self._add_indent(r"\mathbf{function}")
+            + rf" \ \texttt{{{node.name.upper()}}}({', '.join(arg_strs)}) \\ "
+            + body
+            + r" \\ "
+            + self._add_indent(r"\mathbf{end \ function}")
+            + r" \end{array}"
         )
 
     # TODO(ZibingZhang): support \ELSIF
@@ -220,14 +221,14 @@ def visit_If(self, node: ast.If) -> str:
         cond_latex = self._expression_codegen.visit(node.test)
         with self._increment_level():
             body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
-        latex = rf"{self._add_prefix()}\mathbf{{if}} \ {cond_latex} \\ {body_latex}"
+        latex = self._add_indent(rf"\mathbf{{if}} \ {cond_latex} \\ {body_latex}")
 
         if node.orelse:
-            latex += rf" \\ {self._add_prefix()}\mathbf{{else}} \\ "
+            latex += r" \\ " + self._add_indent(r"\mathbf{else} \\ ")
             with self._increment_level():
                 latex += r" \\ ".join(self.visit(stmt) for stmt in node.orelse)
 
-        return latex + rf" \\ {self._add_prefix()}\mathbf{{end \ if}}"
+        return latex + r" \\ " + self._add_indent(r"\mathbf{end \ if}")
 
     def visit_Module(self, node: ast.Module) -> str:
         """Visit a Module node."""
@@ -236,10 +237,10 @@ def visit_Module(self, node: ast.Module) -> str:
     def visit_Return(self, node: ast.Return) -> str:
         """Visit a Return node."""
         return (
-            rf"{self._add_prefix()}\mathbf{{return}}"
-            rf" \ {self._expression_codegen.visit(node.value)}"
+            self._add_indent(r"\mathbf{return}")
+            + rf" \ {self._expression_codegen.visit(node.value)}"
             if node.value is not None
-            else rf"{self._add_prefix()}\mathbf{{return}}"
+            else self._add_indent(r"\mathbf{return}")
         )
 
     def visit_While(self, node: ast.While) -> str:
@@ -253,9 +254,12 @@ def visit_While(self, node: ast.While) -> str:
         with self._increment_level():
             body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
         return (
-            rf"{self._add_prefix()}\mathbf{{while}} \ {cond_latex} \\ "
-            rf"{body_latex} \\ "
-            rf"{self._add_prefix()}\mathbf{{end \ while}}"
+            self._add_indent(r"\mathbf{while} \ ")
+            + cond_latex
+            + r" \\ "
+            + body_latex
+            + r" \\ "
+            + self._add_indent(r"\mathbf{end \ while}")
         )
 
     @contextlib.contextmanager
@@ -265,9 +269,14 @@ def _increment_level(self) -> Generator[None, None, None]:
         yield
         self._indent_level -= 1
 
-    def _add_prefix(self) -> str:
+    def _add_indent(self, line: str) -> str:
+        """Adds an indent before the line.
+
+        Args:
+            line: The line to add an indent to.
+        """
         return (
-            rf"\hspace{{{self._indent_level * self._EM_PER_INDENT}em}} "
+            rf"\hspace{{{self._indent_level * self._EM_PER_INDENT}em}} {line}"
             if self._indent_level > 0
-            else ""
+            else line
         )

From 83bdb7344d253116a2cc8d0416c36e50f1a3004b Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 09:22:56 +0000
Subject: [PATCH 20/27] class arg

---
 src/latexify/codegen/algorithmic_codegen.py | 26 ++++++++++++---------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index a2fdd7b..c73076f 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -153,6 +153,7 @@ class IPythonAlgorithmicCodegen(ast.NodeVisitor):
     """
 
     _EM_PER_INDENT = 1
+    _LINE_BREAK = r" \\ "
 
     _identifier_converter: identifier_converter.IdentifierConverter
     _indent_level: int
@@ -203,14 +204,15 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         # Body
         with self._increment_level():
             body_strs: list[str] = [self.visit(stmt) for stmt in node.body]
-        body = r" \\ ".join(body_strs)
+        body = self._LINE_BREAK.join(body_strs)
 
         return (
             r"\begin{array}{l} "
             + self._add_indent(r"\mathbf{function}")
-            + rf" \ \texttt{{{node.name.upper()}}}({', '.join(arg_strs)}) \\ "
+            + rf" \ \texttt{{{node.name.upper()}}}({', '.join(arg_strs)})"
+            + self._LINE_BREAK
             + body
-            + r" \\ "
+            + self._LINE_BREAK
             + self._add_indent(r"\mathbf{end \ function}")
             + r" \end{array}"
         )
@@ -220,15 +222,17 @@ def visit_If(self, node: ast.If) -> str:
         """Visit an If node."""
         cond_latex = self._expression_codegen.visit(node.test)
         with self._increment_level():
-            body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
-        latex = self._add_indent(rf"\mathbf{{if}} \ {cond_latex} \\ {body_latex}")
+            body_latex = self._LINE_BREAK.join(self.visit(stmt) for stmt in node.body)
+        latex = self._add_indent(
+            rf"\mathbf{{if}} \ {cond_latex}{self._LINE_BREAK}{body_latex}"
+        )
 
         if node.orelse:
-            latex += r" \\ " + self._add_indent(r"\mathbf{else} \\ ")
+            latex += self._LINE_BREAK + self._add_indent(r"\mathbf{else} \\ ")
             with self._increment_level():
-                latex += r" \\ ".join(self.visit(stmt) for stmt in node.orelse)
+                latex += self._LINE_BREAK.join(self.visit(stmt) for stmt in node.orelse)
 
-        return latex + r" \\ " + self._add_indent(r"\mathbf{end \ if}")
+        return latex + self._LINE_BREAK + self._add_indent(r"\mathbf{end \ if}")
 
     def visit_Module(self, node: ast.Module) -> str:
         """Visit a Module node."""
@@ -252,13 +256,13 @@ def visit_While(self, node: ast.While) -> str:
 
         cond_latex = self._expression_codegen.visit(node.test)
         with self._increment_level():
-            body_latex = r" \\ ".join(self.visit(stmt) for stmt in node.body)
+            body_latex = self._LINE_BREAK.join(self.visit(stmt) for stmt in node.body)
         return (
             self._add_indent(r"\mathbf{while} \ ")
             + cond_latex
-            + r" \\ "
+            + self._LINE_BREAK
             + body_latex
-            + r" \\ "
+            + self._LINE_BREAK
             + self._add_indent(r"\mathbf{end \ while}")
         )
 

From b6bafe53050380c98aa5cc2ddde52daa21c3fcf5 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 10:30:26 +0000
Subject: [PATCH 21/27] was weird

---
 src/latexify/codegen/algorithmic_codegen.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index c73076f..89ab048 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -242,7 +242,8 @@ def visit_Return(self, node: ast.Return) -> str:
         """Visit a Return node."""
         return (
             self._add_indent(r"\mathbf{return}")
-            + rf" \ {self._expression_codegen.visit(node.value)}"
+            + r" \ "
+            + self._expression_codegen.visit(node.value)
             if node.value is not None
             else self._add_indent(r"\mathbf{return}")
         )

From 29ca5cb1c84e596dcf8d37a2760e769a05f13db9 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 13:25:19 +0000
Subject: [PATCH 22/27] final (?) suggestions

---
 src/integration_tests/algorithmic_style_test.py  | 4 ++--
 src/latexify/codegen/algorithmic_codegen.py      | 2 +-
 src/latexify/codegen/algorithmic_codegen_test.py | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/integration_tests/algorithmic_style_test.py b/src/integration_tests/algorithmic_style_test.py
index dbe2b5d..8cbe68e 100644
--- a/src/integration_tests/algorithmic_style_test.py
+++ b/src/integration_tests/algorithmic_style_test.py
@@ -29,7 +29,7 @@ def fact(n):
     ).strip()
     ipython_latex = (
         r"\begin{array}{l}"
-        r" \mathbf{function} \ \texttt{FACT}(n) \\"
+        r" \mathbf{function} \ \mathrm{fact}(n) \\"
         r" \hspace{1em} \mathbf{if} \ n = 0 \\"
         r" \hspace{2em} \mathbf{return} \ 1 \\"
         r" \hspace{1em} \mathbf{else} \\"
@@ -74,7 +74,7 @@ def collatz(n):
     ).strip()
     ipython_latex = (
         r"\begin{array}{l}"
-        r" \mathbf{function} \ \texttt{COLLATZ}(n) \\"
+        r" \mathbf{function} \ \mathrm{collatz}(n) \\"
         r" \hspace{1em} \mathrm{iterations} \gets 0 \\"
         r" \hspace{1em} \mathbf{while} \ n > 1 \\"
         r" \hspace{2em} \mathbf{if} \ n \mathbin{\%} 2 = 0 \\"
diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index 89ab048..04a0c4a 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -209,7 +209,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         return (
             r"\begin{array}{l} "
             + self._add_indent(r"\mathbf{function}")
-            + rf" \ \texttt{{{node.name.upper()}}}({', '.join(arg_strs)})"
+            + rf" \ \mathrm{{{node.name}}}({', '.join(arg_strs)})"
             + self._LINE_BREAK
             + body
             + self._LINE_BREAK
diff --git a/src/latexify/codegen/algorithmic_codegen_test.py b/src/latexify/codegen/algorithmic_codegen_test.py
index 58d8244..a0781c7 100644
--- a/src/latexify/codegen/algorithmic_codegen_test.py
+++ b/src/latexify/codegen/algorithmic_codegen_test.py
@@ -188,7 +188,7 @@ def test_visit_assign_jupyter(code: str, latex: str) -> None:
             (
                 r"\begin{array}{l}"
                 r" \mathbf{function}"
-                r" \ \texttt{F}(x) \\"
+                r" \ \mathrm{f}(x) \\"
                 r" \hspace{1em} \mathbf{return} \ x \\"
                 r" \mathbf{end \ function}"
                 r" \end{array}"
@@ -199,7 +199,7 @@ def test_visit_assign_jupyter(code: str, latex: str) -> None:
             (
                 r"\begin{array}{l}"
                 r" \mathbf{function}"
-                r" \ \texttt{F}(a, b, c) \\"
+                r" \ \mathrm{f}(a, b, c) \\"
                 r" \hspace{1em} \mathbf{return} \ 3 \\"
                 r" \mathbf{end \ function}"
                 r" \end{array}"

From 73ecd7313a8d090523ea1b7a869730ac71f1fbeb Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 13:53:19 +0000
Subject: [PATCH 23/27] bad place for f string

---
 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 04a0c4a..08a25b3 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -224,7 +224,7 @@ def visit_If(self, node: ast.If) -> str:
         with self._increment_level():
             body_latex = self._LINE_BREAK.join(self.visit(stmt) for stmt in node.body)
         latex = self._add_indent(
-            rf"\mathbf{{if}} \ {cond_latex}{self._LINE_BREAK}{body_latex}"
+            rf"\mathbf{{if}} \ " + cond_latex + self._LINE_BREAK + body_latex
         )
 
         if node.orelse:

From 55ec2cb82e0f039d7495898ca6cf00f38522c3c2 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 13:53:43 +0000
Subject: [PATCH 24/27] rm f

---
 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 08a25b3..017cdc6 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -224,7 +224,7 @@ def visit_If(self, node: ast.If) -> str:
         with self._increment_level():
             body_latex = self._LINE_BREAK.join(self.visit(stmt) for stmt in node.body)
         latex = self._add_indent(
-            rf"\mathbf{{if}} \ " + cond_latex + self._LINE_BREAK + body_latex
+            r"\mathbf{if} \ " + cond_latex + self._LINE_BREAK + body_latex
         )
 
         if node.orelse:

From 3f0ce8edbdeb245ebb613e47f4c10b20eed61518 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 13:55:29 +0000
Subject: [PATCH 25/27] easier to read

---
 src/integration_tests/integration_utils.py | 12 ++++++------
 src/latexify/ipython_wrappers.py           |  4 ++--
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/integration_tests/integration_utils.py b/src/integration_tests/integration_utils.py
index 5ffe7a0..9f91bbb 100644
--- a/src/integration_tests/integration_utils.py
+++ b/src/integration_tests/integration_utils.py
@@ -26,7 +26,7 @@ def check_function(
     if not kwargs:
         latexified = frontend.function(fn)
         assert str(latexified) == latex
-        assert latexified._repr_latex_() == rf"$$ \displaystyle {latex} $$"
+        assert latexified._repr_latex_() == r"$$ \displaystyle " + latex + " $$"
 
     # Checks the syntax:
     #     @function(**kwargs)
@@ -34,7 +34,7 @@ def check_function(
     #         ...
     latexified = frontend.function(**kwargs)(fn)
     assert str(latexified) == latex
-    assert latexified._repr_latex_() == rf"$$ \displaystyle {latex} $$"
+    assert latexified._repr_latex_() == r"$$ \displaystyle " + latex + " $$"
 
     # Checks the syntax:
     #     def fn(...):
@@ -42,7 +42,7 @@ def check_function(
     #     latexified = function(fn, **kwargs)
     latexified = frontend.function(fn, **kwargs)
     assert str(latexified) == latex
-    assert latexified._repr_latex_() == rf"$$ \displaystyle {latex} $$"
+    assert latexified._repr_latex_() == r"$$ \displaystyle " + latex + " $$"
 
 
 def check_algorithm(
@@ -66,7 +66,7 @@ def check_algorithm(
     if not kwargs:
         latexified = frontend.algorithmic(fn)
         assert str(latexified) == latex
-        assert latexified._repr_latex_() == f"$ {ipython_latex} $"
+        assert latexified._repr_latex_() == "$ " + ipython_latex + " $"
 
     # Checks the syntax:
     #     @algorithmic(**kwargs)
@@ -74,7 +74,7 @@ def check_algorithm(
     #         ...
     latexified = frontend.algorithmic(**kwargs)(fn)
     assert str(latexified) == latex
-    assert latexified._repr_latex_() == f"$ {ipython_latex} $"
+    assert latexified._repr_latex_() == "$ " + ipython_latex + " $"
 
     # Checks the syntax:
     #     def fn(...):
@@ -82,4 +82,4 @@ def check_algorithm(
     #     latexified = algorithmic(fn, **kwargs)
     latexified = frontend.algorithmic(fn, **kwargs)
     assert str(latexified) == latex
-    assert latexified._repr_latex_() == f"$ {ipython_latex} $"
+    assert latexified._repr_latex_() == "$ " + ipython_latex + " $"
diff --git a/src/latexify/ipython_wrappers.py b/src/latexify/ipython_wrappers.py
index c77a869..88aa015 100644
--- a/src/latexify/ipython_wrappers.py
+++ b/src/latexify/ipython_wrappers.py
@@ -95,7 +95,7 @@ def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:
     def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:
         """IPython hook to display LaTeX visualization."""
         return (
-            f"$ {self._ipython_latex} $"
+            r"$ " + self._ipython_latex + " $"
             if self._ipython_latex is not None
             else self._ipython_error
         )
@@ -133,7 +133,7 @@ def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:
     def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:
         """IPython hook to display LaTeX visualization."""
         return (
-            rf"$$ \displaystyle {self._latex} $$"
+            r"$$ \displaystyle " + self._latex + " $$"
             if self._latex is not None
             else self._error
         )

From 452275696eb59ce8d485acd749a4f0adb7fff88f Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 13:59:21 +0000
Subject: [PATCH 26/27] didn't like my code from before

---
 src/latexify/codegen/algorithmic_codegen.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index 017cdc6..402d0ee 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -79,7 +79,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
                 body_strs: list[str] = [self.visit(stmt) for stmt in node.body]
             body_latex = "\n".join(body_strs)
 
-            latex += f"{body_latex}\n"
+            latex += body_latex + "\n"
             latex += self._add_indent("\\EndFunction\n")
         return latex + self._add_indent(r"\end{algorithmic}")
 
@@ -90,10 +90,10 @@ def visit_If(self, node: ast.If) -> str:
         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}")
+        latex = self._add_indent(f"\\If{{${cond_latex}$}}\n" + body_latex)
 
         if node.orelse:
-            latex += "\n" + self._add_indent(r"\Else") + "\n"
+            latex += "\n" + self._add_indent("\\Else\n")
             with self._increment_level():
                 latex += "\n".join(self.visit(stmt) for stmt in node.orelse)
 
@@ -125,7 +125,8 @@ def visit_While(self, node: ast.While) -> str:
             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"
+            + body_latex
+            + "\n"
             + self._add_indent(r"\EndWhile")
         )
 

From 2a47c19b4a6a4b99e260ece736dbf2a10cbb68db Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zibing.zha@gmail.com>
Date: Wed, 14 Dec 2022 15:00:34 +0000
Subject: [PATCH 27/27] no +

---
 src/integration_tests/integration_utils.py  | 12 ++++++------
 src/latexify/codegen/algorithmic_codegen.py | 21 +++++++--------------
 src/latexify/ipython_wrappers.py            |  4 ++--
 3 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/src/integration_tests/integration_utils.py b/src/integration_tests/integration_utils.py
index 9f91bbb..5ffe7a0 100644
--- a/src/integration_tests/integration_utils.py
+++ b/src/integration_tests/integration_utils.py
@@ -26,7 +26,7 @@ def check_function(
     if not kwargs:
         latexified = frontend.function(fn)
         assert str(latexified) == latex
-        assert latexified._repr_latex_() == r"$$ \displaystyle " + latex + " $$"
+        assert latexified._repr_latex_() == rf"$$ \displaystyle {latex} $$"
 
     # Checks the syntax:
     #     @function(**kwargs)
@@ -34,7 +34,7 @@ def check_function(
     #         ...
     latexified = frontend.function(**kwargs)(fn)
     assert str(latexified) == latex
-    assert latexified._repr_latex_() == r"$$ \displaystyle " + latex + " $$"
+    assert latexified._repr_latex_() == rf"$$ \displaystyle {latex} $$"
 
     # Checks the syntax:
     #     def fn(...):
@@ -42,7 +42,7 @@ def check_function(
     #     latexified = function(fn, **kwargs)
     latexified = frontend.function(fn, **kwargs)
     assert str(latexified) == latex
-    assert latexified._repr_latex_() == r"$$ \displaystyle " + latex + " $$"
+    assert latexified._repr_latex_() == rf"$$ \displaystyle {latex} $$"
 
 
 def check_algorithm(
@@ -66,7 +66,7 @@ def check_algorithm(
     if not kwargs:
         latexified = frontend.algorithmic(fn)
         assert str(latexified) == latex
-        assert latexified._repr_latex_() == "$ " + ipython_latex + " $"
+        assert latexified._repr_latex_() == f"$ {ipython_latex} $"
 
     # Checks the syntax:
     #     @algorithmic(**kwargs)
@@ -74,7 +74,7 @@ def check_algorithm(
     #         ...
     latexified = frontend.algorithmic(**kwargs)(fn)
     assert str(latexified) == latex
-    assert latexified._repr_latex_() == "$ " + ipython_latex + " $"
+    assert latexified._repr_latex_() == f"$ {ipython_latex} $"
 
     # Checks the syntax:
     #     def fn(...):
@@ -82,4 +82,4 @@ def check_algorithm(
     #     latexified = algorithmic(fn, **kwargs)
     latexified = frontend.algorithmic(fn, **kwargs)
     assert str(latexified) == latex
-    assert latexified._repr_latex_() == "$ " + ipython_latex + " $"
+    assert latexified._repr_latex_() == f"$ {ipython_latex} $"
diff --git a/src/latexify/codegen/algorithmic_codegen.py b/src/latexify/codegen/algorithmic_codegen.py
index 402d0ee..685663c 100644
--- a/src/latexify/codegen/algorithmic_codegen.py
+++ b/src/latexify/codegen/algorithmic_codegen.py
@@ -79,7 +79,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
                 body_strs: list[str] = [self.visit(stmt) for stmt in node.body]
             body_latex = "\n".join(body_strs)
 
-            latex += body_latex + "\n"
+            latex += f"{body_latex}\n"
             latex += self._add_indent("\\EndFunction\n")
         return latex + self._add_indent(r"\end{algorithmic}")
 
@@ -97,7 +97,7 @@ def visit_If(self, node: ast.If) -> str:
             with self._increment_level():
                 latex += "\n".join(self.visit(stmt) for stmt in node.orelse)
 
-        return latex + "\n" + self._add_indent(r"\EndIf")
+        return f"{latex}\n" + self._add_indent(r"\EndIf")
 
     def visit_Module(self, node: ast.Module) -> str:
         """Visit a Module node."""
@@ -125,8 +125,7 @@ def visit_While(self, node: ast.While) -> str:
             body_latex = "\n".join(self.visit(stmt) for stmt in node.body)
         return (
             self._add_indent(f"\\While{{${cond_latex}$}}\n")
-            + body_latex
-            + "\n"
+            + f"{body_latex}\n"
             + self._add_indent(r"\EndWhile")
         )
 
@@ -211,9 +210,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
             r"\begin{array}{l} "
             + self._add_indent(r"\mathbf{function}")
             + rf" \ \mathrm{{{node.name}}}({', '.join(arg_strs)})"
-            + self._LINE_BREAK
-            + body
-            + self._LINE_BREAK
+            + f"{self._LINE_BREAK}{body}{self._LINE_BREAK}"
             + self._add_indent(r"\mathbf{end \ function}")
             + r" \end{array}"
         )
@@ -225,7 +222,7 @@ def visit_If(self, node: ast.If) -> str:
         with self._increment_level():
             body_latex = self._LINE_BREAK.join(self.visit(stmt) for stmt in node.body)
         latex = self._add_indent(
-            r"\mathbf{if} \ " + cond_latex + self._LINE_BREAK + body_latex
+            rf"\mathbf{{if}} \ {cond_latex}{self._LINE_BREAK}{body_latex}"
         )
 
         if node.orelse:
@@ -242,8 +239,7 @@ def visit_Module(self, node: ast.Module) -> str:
     def visit_Return(self, node: ast.Return) -> str:
         """Visit a Return node."""
         return (
-            self._add_indent(r"\mathbf{return}")
-            + r" \ "
+            self._add_indent(r"\mathbf{return} \ ")
             + self._expression_codegen.visit(node.value)
             if node.value is not None
             else self._add_indent(r"\mathbf{return}")
@@ -261,10 +257,7 @@ def visit_While(self, node: ast.While) -> str:
             body_latex = self._LINE_BREAK.join(self.visit(stmt) for stmt in node.body)
         return (
             self._add_indent(r"\mathbf{while} \ ")
-            + cond_latex
-            + self._LINE_BREAK
-            + body_latex
-            + self._LINE_BREAK
+            + f"{cond_latex}{self._LINE_BREAK}{body_latex}{self._LINE_BREAK}"
             + self._add_indent(r"\mathbf{end \ while}")
         )
 
diff --git a/src/latexify/ipython_wrappers.py b/src/latexify/ipython_wrappers.py
index 88aa015..c77a869 100644
--- a/src/latexify/ipython_wrappers.py
+++ b/src/latexify/ipython_wrappers.py
@@ -95,7 +95,7 @@ def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:
     def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:
         """IPython hook to display LaTeX visualization."""
         return (
-            r"$ " + self._ipython_latex + " $"
+            f"$ {self._ipython_latex} $"
             if self._ipython_latex is not None
             else self._ipython_error
         )
@@ -133,7 +133,7 @@ def _repr_html_(self) -> str | tuple[str, dict[str, Any]] | None:
     def _repr_latex_(self) -> str | tuple[str, dict[str, Any]] | None:
         """IPython hook to display LaTeX visualization."""
         return (
-            r"$$ \displaystyle " + self._latex + " $$"
+            rf"$$ \displaystyle {self._latex} $$"
             if self._latex is not None
             else self._error
         )