Allow docstring in function definition

src/integration_tests/regression_test.py

@@ -279,3 +279,23 @@ 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)
+
+
+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)

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.
 

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

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
@@ -253,6 +253,9 @@ 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):
+                continue
+
             if not isinstance(child, ast.Assign):
                 raise exceptions.LatexifyNotSupportedError(
                     "Codegen supports only Assign nodes in multiline functions, "

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,46 @@ 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 = 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(
     "code,latex",
     [