Support expanding some functions (#125)

* Impl function expander

* Fix error message

* fix flake8 errors

* file doc

* Suggestions

* Fix docstring

* Fix docstring

* rename test methods

* slight rephrasing

* rm local var

* Reformat import

* suggestions

Co-authored-by: Zibing Zhang <zizhang@hubspot.com>

Author: ZibingZhang

src/integration_tests/regression_test.py

@@ -281,6 +281,37 @@ def solve(a, b):
     _check_function(solve, latex)
 
 
+def test_expand_hypot_function_without_attribute_access() -> None:
+    from math import hypot
+
+    def solve(x, y, z):
+        return hypot(x, y, z)
+
+    latex = r"\mathrm{solve}(x, y, z) = \sqrt{x^{{2}} + y^{{2}} + z^{{2}}}"
+    _check_function(solve, latex, expand_functions={"hypot"})
+
+
+def test_expand_hypot_function() -> None:
+    def solve(x, y, z):
+        return math.hypot(x, y, z)
+
+    latex = r"\mathrm{solve}(x, y, z) = \sqrt{x^{{2}} + y^{{2}} + z^{{2}}}"
+    _check_function(solve, latex, expand_functions={"hypot"})
+
+
+def test_expand_nested_function() -> None:
+    def solve(a, b, x, y):
+        return math.hypot(math.hypot(a, b), x, y)
+
+    latex = (
+        r"\mathrm{solve}(a, b, x, y) = "
+        r"\sqrt{"
+        r"\sqrt{a^{{2}} + b^{{2}}}^{{2}} + "
+        r"x^{{2}} + y^{{2}}}"
+    )
+    _check_function(solve, latex, expand_functions={"hypot"})
+
+
 def test_docstring_allowed() -> None:
     def solve(x):
         """The identity function."""

src/latexify/ast_utils.py

@@ -104,3 +104,20 @@ def extract_int(node: ast.expr) -> int:
         raise ValueError(f"Unsupported node to extract int: {type(node).__name__}")
 
     return value
+
+
+def extract_function_name_or_none(node: ast.Call) -> str | None:
+    """Extracts function name from the given Call node.
+
+    Args:
+        node: ast.Call.
+
+    Returns:
+        Extracted function name, or None if not found.
+    """
+    if isinstance(node.func, ast.Name):
+        return node.func.id
+    if isinstance(node.func, ast.Attribute):
+        return node.func.attr
+
+    return None

src/latexify/ast_utils_test.py

@@ -146,3 +146,37 @@ def test_extract_int_invalid() -> None:
         ast_utils.extract_int(ast_utils.make_constant("123"))
     with pytest.raises(ValueError, match=r"^Unsupported node to extract int"):
         ast_utils.extract_int(ast_utils.make_constant(b"123"))
+
+
+@pytest.mark.parametrize(
+    "value,expected",
+    [
+        (
+            ast.Call(
+                func=ast.Name(id="hypot", ctx=ast.Load()),
+                args=[],
+            ),
+            "hypot",
+        ),
+        (
+            ast.Call(
+                func=ast.Attribute(
+                    value=ast.Name(id="math", ctx=ast.Load()),
+                    attr="hypot",
+                    ctx=ast.Load(),
+                ),
+                args=[],
+            ),
+            "hypot",
+        ),
+        (
+            ast.Call(
+                func=ast.Call(func=ast.Name(id="foo", ctx=ast.Load()), args=[]),
+                args=[],
+            ),
+            None,
+        ),
+    ],
+)
+def test_extract_function_name_or_none(value: ast.Call, expected: str | None) -> None:
+    assert ast_utils.extract_function_name_or_none(value) == expected

src/latexify/frontend.py

@@ -13,6 +13,7 @@ def get_latex(
     fn: Callable[..., Any],
     *,
     identifiers: dict[str, str] | None = None,
+    expand_functions: set[str] | None = None,
     reduce_assignments: bool = False,
     use_math_symbols: bool = False,
     use_raw_function_name: bool = False,
@@ -28,6 +29,7 @@ def get_latex(
             the replacements.
             Both keys and values have to represent valid Python identifiers:
             ^[A-Za-z_][A-Za-z0-9_]*$
+        expand_functions: If set, the names of the functions to expand.
         reduce_assignments: If True, assignment statements are used to synthesize
             the final expression.
         use_math_symbols: Whether to convert identifiers with a math symbol surface
@@ -52,6 +54,8 @@ def get_latex(
         tree = transformers.IdentifierReplacer(identifiers).visit(tree)
     if reduce_assignments:
         tree = transformers.AssignmentReducer().visit(tree)
+    if expand_functions is not None:
+        tree = transformers.FunctionExpander(expand_functions).visit(tree)
 
     # Generates LaTeX.
     return codegen.FunctionCodegen(

src/latexify/transformers/__init__.py

@@ -1,6 +1,11 @@
 """Package latexify.transformers."""
 
-from latexify.transformers import assignment_reducer, identifier_replacer
+from latexify.transformers import (
+    assignment_reducer,
+    function_expander,
+    identifier_replacer,
+)
 
 AssignmentReducer = assignment_reducer.AssignmentReducer
+FunctionExpander = function_expander.FunctionExpander
 IdentifierReplacer = identifier_replacer.IdentifierReplacer

src/latexify/transformers/function_expander.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import ast
+import functools
+from collections.abc import Callable
+
+from latexify import ast_utils
+
+
+# TODO(ZibingZhang): handle recursive function expansions
+class FunctionExpander(ast.NodeTransformer):
+    """NodeTransformer to expand functions.
+
+    This class replaces function calls with an expanded form.
+
+    Example:
+        def f(x, y):
+            return hypot(x, y)
+
+        FunctionExpander({"hypot"}) will modify the AST of the function above to below:
+
+        def f(x, y):
+            return sqrt(x**2, y**2)
+    """
+
+    def __init__(self, functions: set[str]) -> None:
+        self._functions = functions
+
+    def visit_Call(self, node: ast.Call) -> ast.AST:
+        """Visitor of Call nodes."""
+        func_name = ast_utils.extract_function_name_or_none(node)
+        if (
+            func_name is not None
+            and func_name in self._functions
+            and func_name in _FUNCTION_EXPANDERS
+        ):
+            return _FUNCTION_EXPANDERS[func_name](self, node)
+
+        return node
+
+
+def _hypot_expander(function_expander: FunctionExpander, node: ast.Call) -> ast.AST:
+    if len(node.args) == 0:
+        return ast_utils.make_constant(0)
+
+    args = [
+        ast.BinOp(function_expander.visit(arg), ast.Pow(), ast_utils.make_constant(2))
+        for arg in node.args
+    ]
+
+    args_reduced = functools.reduce(lambda a, b: ast.BinOp(a, ast.Add(), b), args)
+    return ast.Call(
+        func=ast.Name(id="sqrt", ctx=ast.Load()),
+        args=[args_reduced],
+    )
+
+
+_FUNCTION_EXPANDERS: dict[str, Callable[[FunctionExpander, ast.Call], ast.AST]] = {
+    "hypot": _hypot_expander
+}

src/latexify/transformers/function_expander_test.py

@@ -0,0 +1,97 @@
+"""Tests for latexify.transformers.function_expander."""
+
+from __future__ import annotations
+
+import ast
+import math
+
+from latexify import ast_utils, parser, test_utils
+from latexify.transformers.function_expander import FunctionExpander
+
+
+def _make_ast(args: list[str], body: ast.expr) -> ast.Module:
+    """Helper function to generate an AST for f(x).
+
+    Args:
+        args: The arguments passed to the method.
+        body: The body of the return statement.
+
+    Returns:
+        Generated AST.
+    """
+    return ast.Module(
+        body=[
+            ast.FunctionDef(
+                name="f",
+                args=ast.arguments(
+                    args=[ast.arg(arg=arg) for arg in args],
+                    kwonlyargs=[],
+                    kw_defaults=[],
+                    defaults=[],
+                ),
+                body=[ast.Return(body)],
+                decorator_list=[],
+            )
+        ],
+    )
+
+
+def test_hypot_unchanged_without_attribute_access() -> None:
+    from math import hypot
+
+    def f(x, y):
+        return hypot(x, y)
+
+    expected = _make_ast(
+        ["x", "y"], ast.Call(ast.Name("hypot"), [ast.Name("x"), ast.Name("y")])
+    )
+    transformed = FunctionExpander(set()).visit(parser.parse_function(f))
+    test_utils.assert_ast_equal(transformed, expected)
+
+
+def test_hypot_unchanged() -> None:
+    def f(x, y):
+        return math.hypot(x, y)
+
+    expected = _make_ast(
+        ["x", "y"],
+        ast.Call(
+            ast.Attribute(ast.Name("math"), "hypot", ast.Load()),
+            [ast.Name("x"), ast.Name("y")],
+        ),
+    )
+    transformed = FunctionExpander(set()).visit(parser.parse_function(f))
+    test_utils.assert_ast_equal(transformed, expected)
+
+
+def test_hypot_expanded() -> None:
+    def f(x, y):
+        return math.hypot(x, y)
+
+    expected = _make_ast(
+        ["x", "y"],
+        ast.Call(
+            ast.Name("sqrt"),
+            [
+                ast.BinOp(
+                    ast.BinOp(ast.Name("x"), ast.Pow(), ast_utils.make_constant(2)),
+                    ast.Add(),
+                    ast.BinOp(ast.Name("y"), ast.Pow(), ast_utils.make_constant(2)),
+                )
+            ],
+        ),
+    )
+    transformed = FunctionExpander({"hypot"}).visit(parser.parse_function(f))
+    test_utils.assert_ast_equal(transformed, expected)
+
+
+def test_hypot_expanded_no_args() -> None:
+    def f():
+        return math.hypot()
+
+    expected = _make_ast(
+        [],
+        ast_utils.make_constant(0),
+    )
+    transformed = FunctionExpander({"hypot"}).visit(parser.parse_function(f))
+    test_utils.assert_ast_equal(transformed, expected)