Repository: google/latexify_py

Issue:
Support substitutions of names
According to: https://github.com/google/latexify_py/issues/46#issuecomment-1278292646

Currently, latexify uses only the syntax information of the functions wrapped by `with_latex` and the decorator does not know any bound values outside the function. Some people want to substitute identifiers in the function into its actual value. For example:

```python
# Values bound by literals
t = 42
@with_latex
def diff(y):
    return y - t

# Values bound by given variables
def wrapper(t):
    @with_latex
    def diff(y):
        return y - t

mydiff = wrapper(42)
```

In both cases the user may want to see [ $diff(y) \triangleq y - 42$ ], not [ $diff(y) \triangleq y - t$ ].

I think implementing additional parameter to `with_latex` to specify substitutions, for example:
```python
@with_latex(substitutions={"t": t})
def diff(y):
    ...
```
is a handy solution at this moment, though users are required to write some additional stuff. To implement it, we need to define a substitution rule in the AST visitor that replaces matched identifiers to given strings (or any values that can be converted into a string).

Comments:

Patch:
diff --git a/src/latexify/frontend.py b/src/latexify/frontend.py
index 380a200..5894b82 100644
--- a/src/latexify/frontend.py
+++ b/src/latexify/frontend.py
@@ -11,25 +11,32 @@ from typing import Any
 import dill
 
 from latexify import latexify_visitor
+from latexify.transformers.identifier_replacer import IdentifierReplacer
 
 
 def get_latex(
     fn: Callable[..., Any],
     *,
+    identifiers: dict[str, str] | None = None,
+    reduce_assignments: bool = False,
     use_math_symbols: bool = False,
     use_raw_function_name: bool = False,
-    reduce_assignments: bool = False,
 ) -> str:
     """Obtains LaTeX description from the function's source.
 
     Args:
         fn: Reference to a function to analyze.
+        identifiers: If set, the mapping to replace identifier names in the function.
+            Keys are the original names of the identifiers, and corresponding values are
+            the replacements.
+            Both keys and values have to represent valid Python identifiers:
+            ^[A-Za-z_][A-Za-z0-9_]*$
+        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
             (e.g., "alpha") to the LaTeX symbol (e.g., "\\alpha").
         use_raw_function_name: Whether to keep underscores "_" in the function name,
             or convert it to subscript.
-        reduce_assignments: If True, assignment statements are used to synthesize
-            the final expression.
 
     Returns:
         Generatee LaTeX description.
@@ -45,6 +52,9 @@ def get_latex(
 
     tree = ast.parse(source)
 
+    if identifiers is not None:
+        tree = IdentifierReplacer(identifiers).visit(tree)
+
     visitor = latexify_visitor.LatexifyVisitor(
         use_math_symbols=use_math_symbols,
         use_raw_function_name=use_raw_function_name,
diff --git a/src/latexify/latexify_visitor.py b/src/latexify/latexify_visitor.py
index d271457..476d1ef 100644
--- a/src/latexify/latexify_visitor.py
+++ b/src/latexify/latexify_visitor.py
@@ -23,19 +23,19 @@ class LatexifyVisitor(node_visitor_base.NodeVisitorBase):
     def __init__(
         self,
         *,
+        reduce_assignments: bool = True,
         use_math_symbols: bool = False,
         use_raw_function_name: bool = False,
-        reduce_assignments: bool = True,
     ):
         """Initializer.
 
         Args:
+            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
                 (e.g., "alpha") to the LaTeX symbol (e.g., "\\alpha").
             use_raw_function_name: Whether to keep underscores "_" in the function name,
                 or convert it to subscript.
-            reduce_assignments: If True, assignment statements are used to synthesize
-                the final expression.
         """
         self._math_symbol_converter = math_symbols.MathSymbolConverter(
             enabled=use_math_symbols
diff --git a/src/latexify/transformers/__init__.py b/src/latexify/transformers/__init__.py
new file mode 100644
index 0000000..005a0d2
--- /dev/null
+++ b/src/latexify/transformers/__init__.py
@@ -0,0 +1,1 @@
+"""Package latexify.transformers."""
diff --git a/src/latexify/transformers/identifier_replacer.py b/src/latexify/transformers/identifier_replacer.py
new file mode 100644
index 0000000..3e7f903
--- /dev/null
+++ b/src/latexify/transformers/identifier_replacer.py
@@ -0,0 +1,84 @@
+"""Transformer to replace user symbols."""
+
+from __future__ import annotations
+
+import ast
+import re
+import sys
+from typing import ClassVar
+
+
+class IdentifierReplacer(ast.NodeTransformer):
+    """NodeTransformer to replace identifier names.
+
+    This class defines a rule to replace identifiers in AST with specified names.
+
+    Example:
+        def foo(bar):
+            return baz
+
+        IdentifierReplacer({"foo": "x", "bar": "y", "baz": "z"}) will modify the AST of
+        the function above to below:
+
+        def x(y):
+            return z
+    """
+
+    _IDENTIFIER_PATTERN: ClassVar[re.Pattern] = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+    def __init__(self, mapping: dict[str, str]):
+        """Initializer.
+
+        Args:
+            mapping: User defined mapping of names. Keys are the original names of the
+                identifiers, and corresponding values are the replacements.
+                Both keys and values have to represent valid Python identifiers:
+                ^[A-Za-z_][A-Za-z0-9_]*$
+        """
+        self._mapping = mapping
+
+        for k, v in self._mapping.items():
+            if not self._IDENTIFIER_PATTERN.match(k):
+                raise ValueError(f"'{k}' is not an identifier name.")
+            if not self._IDENTIFIER_PATTERN.match(v):
+                raise ValueError(f"'{v}' is not an identifier name.")
+
+    def _replace_args(self, args: list[ast.arg]) -> list[ast.arg]:
+        """Helper function to replace arg names."""
+        return [ast.arg(arg=self._mapping.get(a.arg, a.arg)) for a in args]
+
+    def _visit_children(self, children: list[ast.AST]) -> list[ast.AST]:
+        """Helper function to visit all children."""
+        return [self.visit(child) for child in children]
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
+        """Visitor of FunctionDef."""
+        if sys.version_info.minor < 8:
+            args = ast.arguments(
+                args=self._replace_args(node.args.args),
+                kwonlyargs=self._replace_args(node.args.kwonlyargs),
+                kw_defaults=self._visit_children(node.args.kw_defaults),
+                defaults=self._visit_children(node.args.defaults),
+            )
+        else:
+            args = ast.arguments(
+                posonlyargs=self._replace_args(node.args.posonlyargs),  # from 3.8
+                args=self._replace_args(node.args.args),
+                kwonlyargs=self._replace_args(node.args.kwonlyargs),
+                kw_defaults=self._visit_children(node.args.kw_defaults),
+                defaults=self._visit_children(node.args.defaults),
+            )
+
+        return ast.FunctionDef(
+            name=self._mapping.get(node.name, node.name),
+            args=args,
+            body=self._visit_children(node.body),
+            decorator_list=self._visit_children(node.decorator_list),
+        )
+
+    def visit_Name(self, node: ast.Name) -> ast.Name:
+        """Visitor of Name."""
+        return ast.Name(
+            id=self._mapping.get(node.id, node.id),
+            ctx=node.ctx,
+        )


Test patch:
diff --git a/src/latexify/frontend_test.py b/src/latexify/frontend_test.py
new file mode 100644
index 0000000..91a9d30
--- /dev/null
+++ b/src/latexify/frontend_test.py
@@ -0,0 +1,19 @@
+"""Tests for latexify.frontend."""
+
+from __future__ import annotations
+
+from latexify import frontend
+
+
+def test_identifiers() -> None:
+    def very_long_name_function(very_long_name_variable):
+        return 3 * very_long_name_variable
+
+    identifiers = {
+        "very_long_name_function": "f",
+        "very_long_name_variable": "x",
+    }
+
+    assert frontend.get_latex(very_long_name_function, identifiers=identifiers) == (
+        r"\mathrm{f}(x) \triangleq 3x"
+    )
diff --git a/src/latexify/test_utils.py b/src/latexify/test_utils.py
new file mode 100644
index 0000000..b45da9b
--- /dev/null
+++ b/src/latexify/test_utils.py
@@ -0,0 +1,59 @@
+"""Test utilities."""
+
+from __future__ import annotations
+
+import ast
+from typing import cast
+
+
+def ast_equal(tree1: ast.AST, tree2: ast.AST) -> bool:
+    """Checks the equality between two ASTs.
+
+    Args:
+        tree1: An AST to compare.
+        tree2: Another AST.
+
+    Returns:
+        True if tree1 and tree2 represent the same AST, False otherwise.
+    """
+    try:
+        assert type(tree1) is type(tree2)
+
+        for k, v1 in vars(tree1).items():
+            v2 = getattr(tree2, k)
+
+            if isinstance(v1, ast.AST):
+                assert ast_equal(v1, cast(ast.AST, v2))
+            elif isinstance(v1, list):
+                v2 = cast(list, v2)
+                assert len(v1) == len(v2)
+                assert all(
+                    ast_equal(cast(ast.AST, c1), cast(ast.AST, c2))
+                    for c1, c2 in zip(v1, v2)
+                )
+            else:
+                assert v1 == v2
+
+    except AssertionError:
+        return False
+
+    return True
+
+
+def assert_ast_equal(tree1: ast.AST, tree2: ast.AST) -> None:
+    """Asserts the equality between two ASTs.
+
+    Args:
+        tree1: An AST to compare.
+        tree2: Another AST.
+
+    Raises:
+        AssertionError: tree1 and tree2 represent different ASTs.
+    """
+    assert ast_equal(
+        tree1, tree2
+    ), f"""\
+AST does not match.
+tree1={ast.dump(tree1, indent=4)}
+tree2={ast.dump(tree2, indent=4)}
+"""
diff --git a/src/latexify/transformers/identifier_replacer_test.py b/src/latexify/transformers/identifier_replacer_test.py
new file mode 100644
index 0000000..f555b92
--- /dev/null
+++ b/src/latexify/transformers/identifier_replacer_test.py
@@ -0,0 +1,101 @@
+"""Tests for latexify.transformer.identifier_replacer."""
+
+import ast
+
+import pytest
+
+from latexify import test_utils
+from latexify.transformers.identifier_replacer import IdentifierReplacer
+
+
+def test_invalid_mapping() -> None:
+    with pytest.raises(ValueError, match=r"'123' is not an identifier name."):
+        IdentifierReplacer({"123": "foo"})
+    with pytest.raises(ValueError, match=r"'456' is not an identifier name."):
+        IdentifierReplacer({"foo": "456"})
+
+
+def test_name_replaced() -> None:
+    source = ast.Name(id="foo", ctx=ast.Load())
+    expected = ast.Name(id="bar", ctx=ast.Load())
+    transformed = IdentifierReplacer({"foo": "bar"}).visit(source)
+    test_utils.assert_ast_equal(transformed, expected)
+
+
+def test_name_not_replaced() -> None:
+    source = ast.Name(id="foo", ctx=ast.Load())
+    expected = ast.Name(id="foo", ctx=ast.Load())
+    transformed = IdentifierReplacer({"fo": "bar"}).visit(source)
+    test_utils.assert_ast_equal(transformed, expected)
+    transformed = IdentifierReplacer({"fooo": "bar"}).visit(source)
+    test_utils.assert_ast_equal(transformed, expected)
+
+
+def test_functiondef() -> None:
+    # Subtree of:
+    #     @d
+    #     def f(x=a, /, y=b, *, z=c):
+    #         ...
+    source = ast.FunctionDef(
+        name="f",
+        args=ast.arguments(
+            posonlyargs=[ast.arg(arg="x")],
+            args=[ast.arg(arg="y")],
+            kwonlyargs=[ast.arg(arg="z")],
+            kw_defaults=[ast.Name(id="c", ctx=ast.Load())],
+            defaults=[
+                ast.Name(id="a", ctx=ast.Load()),
+                ast.Name(id="b", ctx=ast.Load()),
+            ],
+        ),
+        body=[ast.Ellipsis()],
+        decorator_list=[ast.Name(id="d", ctx=ast.Load())],
+    )
+
+    expected = ast.FunctionDef(
+        name="F",
+        args=ast.arguments(
+            posonlyargs=[ast.arg(arg="X")],
+            args=[ast.arg(arg="Y")],
+            kwonlyargs=[ast.arg(arg="Z")],
+            kw_defaults=[ast.Name(id="C", ctx=ast.Load())],
+            defaults=[
+                ast.Name(id="A", ctx=ast.Load()),
+                ast.Name(id="B", ctx=ast.Load()),
+            ],
+        ),
+        body=[ast.Ellipsis()],
+        decorator_list=[ast.Name(id="D", ctx=ast.Load())],
+    )
+
+    mapping = {x: x.upper() for x in "abcdfxyz"}
+    transformed = IdentifierReplacer(mapping).visit(source)
+    test_utils.assert_ast_equal(transformed, expected)
+
+
+def test_expr() -> None:
+    # Subtree of:
+    #     (x + y) * z
+    source = ast.BinOp(
+        left=ast.BinOp(
+            left=ast.Name(id="x", ctx=ast.Load()),
+            op=ast.Add(),
+            right=ast.Name(id="y", ctx=ast.Load()),
+        ),
+        op=ast.Mult(),
+        right=ast.Name(id="z", ctx=ast.Load()),
+    )
+
+    expected = ast.BinOp(
+        left=ast.BinOp(
+            left=ast.Name(id="X", ctx=ast.Load()),
+            op=ast.Add(),
+            right=ast.Name(id="Y", ctx=ast.Load()),
+        ),
+        op=ast.Mult(),
+        right=ast.Name(id="Z", ctx=ast.Load()),
+    )
+
+    mapping = {x: x.upper() for x in "xyz"}
+    transformed = IdentifierReplacer(mapping).visit(source)
+    test_utils.assert_ast_equal(transformed, expected)