From e05181275a2cca210cf0f6ad74df58634c681bfe Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sat, 19 Nov 2022 23:58:26 +0000
Subject: [PATCH 1/7] Allow docstring

---
 src/integration_tests/regression_test.py      |  9 ++++++++
 src/latexify/codegen/function_codegen.py      | 10 ++++++++-
 src/latexify/codegen/function_codegen_test.py | 21 ++++++++++++++++++-
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/integration_tests/regression_test.py b/src/integration_tests/regression_test.py
index 89c59e0..0c30ad3 100644
--- a/src/integration_tests/regression_test.py
+++ b/src/integration_tests/regression_test.py
@@ -279,3 +279,12 @@ def solve(a, b):
         r"a - b \mathclose{}\right) - a b"
     )
     _check_function(solve, latex)
+
+
+def test_docstring_allowed() -> None:
+    def solve(x):
+        """The identity function."""
+        return x
+
+    latex = r"\mathrm{solve}(x) = x"
+    _check_function(solve, latex)
diff --git a/src/latexify/codegen/function_codegen.py b/src/latexify/codegen/function_codegen.py
index f25c9ad..135f387 100644
--- a/src/latexify/codegen/function_codegen.py
+++ b/src/latexify/codegen/function_codegen.py
@@ -252,7 +252,15 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         body_strs: list[str] = []
 
         # Assignment statements (if any): x = ...
-        for child in node.body[:-1]:
+        for i, child in enumerate(node.body[:-1]):
+            # Allow docstrings
+            if (
+                i == 0
+                and isinstance(child, ast.Expr)
+                and isinstance(child.value, ast.Constant)
+            ):
+                continue
+
             if not isinstance(child, ast.Assign):
                 raise exceptions.LatexifyNotSupportedError(
                     "Codegen supports only Assign nodes in multiline functions, "
diff --git a/src/latexify/codegen/function_codegen_test.py b/src/latexify/codegen/function_codegen_test.py
index 30f4436..8f5006a 100644
--- a/src/latexify/codegen/function_codegen_test.py
+++ b/src/latexify/codegen/function_codegen_test.py
@@ -6,7 +6,7 @@
 
 import pytest
 
-from latexify import exceptions, test_utils
+from latexify import ast_utils, exceptions, test_utils
 from latexify.codegen import FunctionCodegen, function_codegen
 
 
@@ -40,6 +40,25 @@ def test_visit_functiondef_use_signature() -> None:
     assert FunctionCodegen(use_signature=True).visit(tree) == latex_with_flag
 
 
+def test_visit_functiondef_ignore_docstring() -> None:
+    tree = ast.FunctionDef(
+        name="f",
+        args=ast.arguments(
+            args=[ast.arg(arg="x")],
+            kwonlyargs=[],
+            kw_defaults=[],
+            defaults=[],
+        ),
+        body=[
+            ast.Expr(ast_utils.make_constant("""docstring""")),
+            ast.Return(value=ast.Name(id="x", ctx=ast.Load())),
+        ],
+        decorator_list=[],
+    )
+    latex_with_flag = r"\mathrm{f}(x) = x"
+    assert FunctionCodegen().visit(tree) == latex_with_flag
+
+
 @pytest.mark.parametrize(
     "code,latex",
     [

From 66c9f74b7ca59083fa0684f0b6dd14e59c0a5ccf Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 20 Nov 2022 08:01:16 +0000
Subject: [PATCH 2/7] Suggestions

---
 src/integration_tests/regression_test.py      | 11 +++++++
 src/latexify/ast_utils.py                     | 25 +++++++++++++++
 src/latexify/ast_utils_test.py                | 31 +++++++++++++++++++
 src/latexify/codegen/function_codegen.py      | 12 +++----
 src/latexify/codegen/function_codegen_test.py | 27 ++++++++++++++--
 5 files changed, 95 insertions(+), 11 deletions(-)

diff --git a/src/integration_tests/regression_test.py b/src/integration_tests/regression_test.py
index 0c30ad3..19ceac9 100644
--- a/src/integration_tests/regression_test.py
+++ b/src/integration_tests/regression_test.py
@@ -288,3 +288,14 @@ def solve(x):
 
     latex = r"\mathrm{solve}(x) = x"
     _check_function(solve, latex)
+
+
+def test_multiple_constants_allowed() -> None:
+    def solve(x):
+        """The identity function."""
+        123
+        True
+        return x
+
+    latex = r"\mathrm{solve}(x) = x"
+    _check_function(solve, latex)
diff --git a/src/latexify/ast_utils.py b/src/latexify/ast_utils.py
index 2d2e6c9..8a686de 100644
--- a/src/latexify/ast_utils.py
+++ b/src/latexify/ast_utils.py
@@ -41,6 +41,31 @@ def make_constant(value: Any) -> ast.expr:
     raise ValueError(f"Unsupported type to generate Constant: {type(value).__name__}")
 
 
+def is_constant(node: ast.AST) -> bool:
+    """Checks if the node is a constant.
+
+    Args:
+        node: The node to examine.
+
+    Returns:
+        True if the node is a constant, False otherwise.
+    """
+    if sys.version_info.minor < 8:
+        if isinstance(
+            node,
+            (ast.Bytes, ast.Constant, ast.Ellipsis, ast.NameConstant, ast.Num, ast.Str),
+        ):
+            return True
+    else:
+        if isinstance(node, ast.Constant):
+            return True
+
+    if isinstance(node, ast.Expr):
+        return is_constant(node.value)
+
+    return False
+
+
 def extract_int_or_none(node: ast.expr) -> int | None:
     """Extracts int constant from the given Constant node.
 
diff --git a/src/latexify/ast_utils_test.py b/src/latexify/ast_utils_test.py
index ab3352a..00da0e1 100644
--- a/src/latexify/ast_utils_test.py
+++ b/src/latexify/ast_utils_test.py
@@ -59,6 +59,37 @@ def test_make_constant_invalid() -> None:
         ast_utils.make_constant(object())
 
 
+@test_utils.require_at_most(7)
+@pytest.mark.parametrize(
+    "value,expected",
+    [
+        (ast.Bytes(s=b"foo"), True),
+        (ast.Constant("bar"), True),
+        (ast.Ellipsis(), True),
+        (ast.NameConstant(value=None), True),
+        (ast.Num(n=123), True),
+        (ast.Str(s="baz"), True),
+        (ast.Expr(value=ast.Num(456)), True),
+        (ast.Global("qux"), False),
+    ],
+)
+def test_is_constant_legacy(value: Any, expected: ast.Constant) -> None:
+    assert ast_utils.is_constant(value) is expected
+
+
+@test_utils.require_at_least(8)
+@pytest.mark.parametrize(
+    "value,expected",
+    [
+        (ast.Constant("foo"), True),
+        (ast.Expr(value=ast.Constant(123)), True),
+        (ast.Global("bar"), False),
+    ],
+)
+def test_is_constant(value: Any, expected: ast.Constant) -> None:
+    assert ast_utils.is_constant(value) is expected
+
+
 def test_extract_int_or_none() -> None:
     assert ast_utils.extract_int_or_none(ast_utils.make_constant(-123)) == -123
     assert ast_utils.extract_int_or_none(ast_utils.make_constant(0)) == 0
diff --git a/src/latexify/codegen/function_codegen.py b/src/latexify/codegen/function_codegen.py
index 135f387..e96a5fc 100644
--- a/src/latexify/codegen/function_codegen.py
+++ b/src/latexify/codegen/function_codegen.py
@@ -7,7 +7,7 @@
 import sys
 from typing import Any
 
-from latexify import analyzers, constants, exceptions, math_symbols
+from latexify import analyzers, ast_utils, constants, exceptions, math_symbols
 
 # Precedences of operators for BoolOp, BinOp, UnaryOp, and Compare nodes.
 # Note that this value affects only the appearance of surrounding parentheses for each
@@ -252,13 +252,9 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
         body_strs: list[str] = []
 
         # Assignment statements (if any): x = ...
-        for i, child in enumerate(node.body[:-1]):
-            # Allow docstrings
-            if (
-                i == 0
-                and isinstance(child, ast.Expr)
-                and isinstance(child.value, ast.Constant)
-            ):
+        for child in node.body[:-1]:
+            # Allow constants
+            if ast_utils.is_constant(child):
                 continue
 
             if not isinstance(child, ast.Assign):
diff --git a/src/latexify/codegen/function_codegen_test.py b/src/latexify/codegen/function_codegen_test.py
index 8f5006a..ac1e3bc 100644
--- a/src/latexify/codegen/function_codegen_test.py
+++ b/src/latexify/codegen/function_codegen_test.py
@@ -50,13 +50,34 @@ def test_visit_functiondef_ignore_docstring() -> None:
             defaults=[],
         ),
         body=[
-            ast.Expr(ast_utils.make_constant("""docstring""")),
+            ast.Expr(ast_utils.make_constant("docstring")),
             ast.Return(value=ast.Name(id="x", ctx=ast.Load())),
         ],
         decorator_list=[],
     )
-    latex_with_flag = r"\mathrm{f}(x) = x"
-    assert FunctionCodegen().visit(tree) == latex_with_flag
+    latex = r"\mathrm{f}(x) = x"
+    assert FunctionCodegen().visit(tree) == latex
+
+
+def test_visit_functiondef_ignore_multiple_constants() -> None:
+    tree = ast.FunctionDef(
+        name="f",
+        args=ast.arguments(
+            args=[ast.arg(arg="x")],
+            kwonlyargs=[],
+            kw_defaults=[],
+            defaults=[],
+        ),
+        body=[
+            ast.Expr(ast_utils.make_constant("docstring" "")),
+            ast.Expr(ast_utils.make_constant(3)),
+            ast.Expr(ast_utils.make_constant(True)),
+            ast.Return(value=ast.Name(id="x", ctx=ast.Load())),
+        ],
+        decorator_list=[],
+    )
+    latex = r"\mathrm{f}(x) = x"
+    assert FunctionCodegen().visit(tree) == latex
 
 
 @pytest.mark.parametrize(

From 398b660e5241eaddd66458a3da4c9e29e6368d77 Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 20 Nov 2022 08:08:32 +0000
Subject: [PATCH 3/7] Remove obvious comment

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

diff --git a/src/latexify/codegen/function_codegen.py b/src/latexify/codegen/function_codegen.py
index e96a5fc..b73e104 100644
--- a/src/latexify/codegen/function_codegen.py
+++ b/src/latexify/codegen/function_codegen.py
@@ -253,7 +253,6 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
 
         # Assignment statements (if any): x = ...
         for child in node.body[:-1]:
-            # Allow constants
             if ast_utils.is_constant(child):
                 continue
 

From 713be26d63cd3f5e30fde807b7450847f3c0471b Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 20 Nov 2022 08:41:03 +0000
Subject: [PATCH 4/7] use ast.parse

---
 src/latexify/codegen/function_codegen_test.py | 57 ++++++-------------
 1 file changed, 18 insertions(+), 39 deletions(-)

diff --git a/src/latexify/codegen/function_codegen_test.py b/src/latexify/codegen/function_codegen_test.py
index ac1e3bc..bde9317 100644
--- a/src/latexify/codegen/function_codegen_test.py
+++ b/src/latexify/codegen/function_codegen_test.py
@@ -6,7 +6,7 @@
 
 import pytest
 
-from latexify import ast_utils, exceptions, test_utils
+from latexify import exceptions, test_utils
 from latexify.codegen import FunctionCodegen, function_codegen
 
 
@@ -22,16 +22,11 @@ class UnknownNode(ast.AST):
 
 
 def test_visit_functiondef_use_signature() -> None:
-    tree = ast.FunctionDef(
-        name="f",
-        args=ast.arguments(
-            args=[ast.arg(arg="x")],
-            kwonlyargs=[],
-            kw_defaults=[],
-            defaults=[],
-        ),
-        body=[ast.Return(value=ast.Name(id="x", ctx=ast.Load()))],
-        decorator_list=[],
+    tree = ast.parse(
+        """
+def f(x):
+    return x
+    """
     )
     latex_without_flag = "x"
     latex_with_flag = r"\mathrm{f}(x) = x"
@@ -41,40 +36,24 @@ def test_visit_functiondef_use_signature() -> None:
 
 
 def test_visit_functiondef_ignore_docstring() -> None:
-    tree = ast.FunctionDef(
-        name="f",
-        args=ast.arguments(
-            args=[ast.arg(arg="x")],
-            kwonlyargs=[],
-            kw_defaults=[],
-            defaults=[],
-        ),
-        body=[
-            ast.Expr(ast_utils.make_constant("docstring")),
-            ast.Return(value=ast.Name(id="x", ctx=ast.Load())),
-        ],
-        decorator_list=[],
+    tree = ast.parse(
+        """
+def f(x):
+    '''docstring'''
+    return x"""
     )
     latex = r"\mathrm{f}(x) = x"
     assert FunctionCodegen().visit(tree) == latex
 
 
 def test_visit_functiondef_ignore_multiple_constants() -> None:
-    tree = ast.FunctionDef(
-        name="f",
-        args=ast.arguments(
-            args=[ast.arg(arg="x")],
-            kwonlyargs=[],
-            kw_defaults=[],
-            defaults=[],
-        ),
-        body=[
-            ast.Expr(ast_utils.make_constant("docstring" "")),
-            ast.Expr(ast_utils.make_constant(3)),
-            ast.Expr(ast_utils.make_constant(True)),
-            ast.Return(value=ast.Name(id="x", ctx=ast.Load())),
-        ],
-        decorator_list=[],
+    tree = ast.parse(
+        """
+def f(x):
+    '''docstring'''
+    3
+    True
+    return x"""
     )
     latex = r"\mathrm{f}(x) = x"
     assert FunctionCodegen().visit(tree) == latex

From ffa04c66fc6c12717169ee7cdd68d6a4f80c12fa Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 20 Nov 2022 08:46:41 +0000
Subject: [PATCH 5/7] edited suggested format

---
 src/latexify/codegen/function_codegen_test.py | 31 ++++++++++---------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/src/latexify/codegen/function_codegen_test.py b/src/latexify/codegen/function_codegen_test.py
index bde9317..acaf3c8 100644
--- a/src/latexify/codegen/function_codegen_test.py
+++ b/src/latexify/codegen/function_codegen_test.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import ast
+import textwrap
 
 import pytest
 
@@ -36,25 +37,27 @@ def f(x):
 
 
 def test_visit_functiondef_ignore_docstring() -> None:
-    tree = ast.parse(
-        """
-def f(x):
-    '''docstring'''
-    return x"""
-    )
+    tree = ast.parse(textwrap.dedent("""
+        def f(x):
+            '''docstring'''
+            return x
+        """)).body[0]
+    assert isinstance(tree, ast.FunctionDef)
+
     latex = r"\mathrm{f}(x) = x"
     assert FunctionCodegen().visit(tree) == latex
 
 
 def test_visit_functiondef_ignore_multiple_constants() -> None:
-    tree = ast.parse(
-        """
-def f(x):
-    '''docstring'''
-    3
-    True
-    return x"""
-    )
+    tree = ast.parse(textwrap.dedent("""
+        def f(x):
+            '''docstring'''
+            3
+            True
+            return x
+        """)).body[0]
+    assert isinstance(tree, ast.FunctionDef)
+
     latex = r"\mathrm{f}(x) = x"
     assert FunctionCodegen().visit(tree) == latex
 

From 16b564a565b82f501ba59c287ed9f78d46cea80e Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 20 Nov 2022 08:54:03 +0000
Subject: [PATCH 6/7] Suggestions 2

---
 src/latexify/ast_utils.py                     | 13 +++----------
 src/latexify/ast_utils_test.py                |  8 ++++----
 src/latexify/codegen/function_codegen.py      |  2 +-
 src/latexify/codegen/function_codegen_test.py | 16 ++++++++++++----
 4 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/src/latexify/ast_utils.py b/src/latexify/ast_utils.py
index 8a686de..56de6d1 100644
--- a/src/latexify/ast_utils.py
+++ b/src/latexify/ast_utils.py
@@ -51,19 +51,12 @@ def is_constant(node: ast.AST) -> bool:
         True if the node is a constant, False otherwise.
     """
     if sys.version_info.minor < 8:
-        if isinstance(
+        return isinstance(
             node,
             (ast.Bytes, ast.Constant, ast.Ellipsis, ast.NameConstant, ast.Num, ast.Str),
-        ):
-            return True
+        )
     else:
-        if isinstance(node, ast.Constant):
-            return True
-
-    if isinstance(node, ast.Expr):
-        return is_constant(node.value)
-
-    return False
+        return isinstance(node, ast.Constant)
 
 
 def extract_int_or_none(node: ast.expr) -> int | None:
diff --git a/src/latexify/ast_utils_test.py b/src/latexify/ast_utils_test.py
index 00da0e1..bcf0e5b 100644
--- a/src/latexify/ast_utils_test.py
+++ b/src/latexify/ast_utils_test.py
@@ -69,11 +69,11 @@ def test_make_constant_invalid() -> None:
         (ast.NameConstant(value=None), True),
         (ast.Num(n=123), True),
         (ast.Str(s="baz"), True),
-        (ast.Expr(value=ast.Num(456)), True),
+        (ast.Expr(value=ast.Num(456)), False),
         (ast.Global("qux"), False),
     ],
 )
-def test_is_constant_legacy(value: Any, expected: ast.Constant) -> None:
+def test_is_constant_legacy(value: ast.AST, expected: bool) -> None:
     assert ast_utils.is_constant(value) is expected
 
 
@@ -82,11 +82,11 @@ def test_is_constant_legacy(value: Any, expected: ast.Constant) -> None:
     "value,expected",
     [
         (ast.Constant("foo"), True),
-        (ast.Expr(value=ast.Constant(123)), True),
+        (ast.Expr(value=ast.Constant(123)), False),
         (ast.Global("bar"), False),
     ],
 )
-def test_is_constant(value: Any, expected: ast.Constant) -> None:
+def test_is_constant(value: ast.AST, expected: bool) -> None:
     assert ast_utils.is_constant(value) is expected
 
 
diff --git a/src/latexify/codegen/function_codegen.py b/src/latexify/codegen/function_codegen.py
index b73e104..33f8df0 100644
--- a/src/latexify/codegen/function_codegen.py
+++ b/src/latexify/codegen/function_codegen.py
@@ -253,7 +253,7 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> str:
 
         # Assignment statements (if any): x = ...
         for child in node.body[:-1]:
-            if ast_utils.is_constant(child):
+            if isinstance(child, ast.Expr) and ast_utils.is_constant(child.value):
                 continue
 
             if not isinstance(child, ast.Assign):
diff --git a/src/latexify/codegen/function_codegen_test.py b/src/latexify/codegen/function_codegen_test.py
index acaf3c8..9d7947e 100644
--- a/src/latexify/codegen/function_codegen_test.py
+++ b/src/latexify/codegen/function_codegen_test.py
@@ -37,11 +37,15 @@ def f(x):
 
 
 def test_visit_functiondef_ignore_docstring() -> None:
-    tree = ast.parse(textwrap.dedent("""
+    tree = ast.parse(
+        textwrap.dedent(
+            """
         def f(x):
             '''docstring'''
             return x
-        """)).body[0]
+        """
+        )
+    ).body[0]
     assert isinstance(tree, ast.FunctionDef)
 
     latex = r"\mathrm{f}(x) = x"
@@ -49,13 +53,17 @@ def f(x):
 
 
 def test_visit_functiondef_ignore_multiple_constants() -> None:
-    tree = ast.parse(textwrap.dedent("""
+    tree = ast.parse(
+        textwrap.dedent(
+            """
         def f(x):
             '''docstring'''
             3
             True
             return x
-        """)).body[0]
+            """
+        )
+    ).body[0]
     assert isinstance(tree, ast.FunctionDef)
 
     latex = r"\mathrm{f}(x) = x"

From d3eeb754d4633674020f6bd8a8bf8b533bcd314e Mon Sep 17 00:00:00 2001
From: Zibing Zhang <zizhang@hubspot.com>
Date: Sun, 20 Nov 2022 08:56:58 +0000
Subject: [PATCH 7/7] formatting

---
 src/latexify/codegen/function_codegen_test.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/latexify/codegen/function_codegen_test.py b/src/latexify/codegen/function_codegen_test.py
index 9d7947e..793ab62 100644
--- a/src/latexify/codegen/function_codegen_test.py
+++ b/src/latexify/codegen/function_codegen_test.py
@@ -24,11 +24,15 @@ class UnknownNode(ast.AST):
 
 def test_visit_functiondef_use_signature() -> None:
     tree = ast.parse(
-        """
-def f(x):
-    return x
-    """
-    )
+        textwrap.dedent(
+            """
+        def f(x):
+            return x
+            """
+        )
+    ).body[0]
+    assert isinstance(tree, ast.FunctionDef)
+
     latex_without_flag = "x"
     latex_with_flag = r"\mathrm{f}(x) = x"
     assert FunctionCodegen().visit(tree) == latex_with_flag
@@ -43,7 +47,7 @@ def test_visit_functiondef_ignore_docstring() -> None:
         def f(x):
             '''docstring'''
             return x
-        """
+            """
         )
     ).body[0]
     assert isinstance(tree, ast.FunctionDef)