Support substitutions of names about latexify_py HOT 12 CLOSED

But Attribute is hard to replace, doing sub-tree matching feels too complicated.

For now we can abandon this support. Only replacing long Names is still useful.

do str.replace after inspect.getsource

This may break parsing very frequently; I think it would be better not to provide this option.

Or we can annotation inside function.

This is very interesting idea, but it seems to require changing the syntax of the target function by users themselves. Since the significant advantage of this library is generating LaTeX without modifying the original function, it may drop some usability if we started to require such modification.

That being said, I'd like to introduce an additional transformation between parse and LatexifyVisitor as we discussed above.

Currently I have certain time to work on this. Please let me think the actual interface.
(Though I think #50 (comment) can be used as the base implementation)

from latexify_py.

I think I need something like this, but for a simpler reason:

@with_latex(substitutions={"my_long_function_name": "f", "alpha_t_minus_1": "{\alpha}_{t-1}"})
def my_long_function_name(alpha_t_minus_1):
     return alpha_t_minus_1

from latexify_py.

@cccntu

"alpha_t_minus_1": "{\alpha}_{t-1}"

Yes I intended it: the substitutions argument takes the exact representation of a specific identifier.

"my_long_function_name": "f"

This requires several consideration as function names are treated separately. Would providing another argument for this purpose be better?

# name takes the LaTeX form of the function name to substitute,
# while identifiers takes that of inner identifiers.
@with_latex(name="f", identifiers={"alpha_t_minus_1": "{\alpha}_{t-1}"})
def my_long_function_name(alpha_t_minus_1):
     return alpha_t_minus_1

from latexify_py.

That's ok. It would be great if the function names can be substituted, too.

@with_latex(function_name="\epsilon_{\theta}",
 identifiers={"alpha_t_minus_1": "{\alpha}_{t-1}", "torch.sqrt":(r"\sqrt{", "}")})
def my_long_function_name(alpha_t_minus_1):
     return torch.sqrt(alpha_t_minus_1)

Although maybe it's easier to just do this:

sqrt = torch.sqrt
@with_latex(function_name="\epsilon_{\theta}",
 identifiers={"alpha_t_minus_1": "{\alpha}_{t-1}"})
def my_long_function_name(alpha_t_minus_1):
     return sqrt(alpha_t_minus_1)

from latexify_py.

@odashi
I was curious how ast works. So I tried to implement it. Here is a very simple version that seems to work.

import ast
import inspect
from ast import AST, Constant, Load, Name, NodeTransformer, Subscript, fix_missing_locations, iter_fields


class MyTransformer(ast.NodeTransformer):
    def __init__(self, identifiers=None):
        self.identifiers = {} if identifiers is None else identifiers

    def visit(self, node):
        for field, value in iter_fields(node):
            if isinstance(value, str):
                value = self.identifiers.get(value, value)
                setattr(node, field, value)
        return super().visit(node)

tree = ast.parse('foo.bar', mode='eval')

new_tree = fix_missing_locations(MyTransformer({'foo': 'FOO'}).visit(tree))
print(ast.dump(new_tree, indent=4))

output:

Expression(
    body=Attribute(
        value=Name(id='FOO', ctx=Load()),
        attr='bar',
        ctx=Load()))

from latexify_py.

@cccntu Great. Actually we are using the modified version of it (in node_visitor_base.py in this repository) because of some demands of the earlier implementation, but I think it is not so sophisticated nowadays. It may be good opportunity to rethink the overall parsing strategy of this library.
(Since it is not a lightweight subtask, it may also be good to focus on implementing the original suggestion in this issue on the current implementation)

from latexify_py.

@odashi
I'm not sure what you have in mind.
But it's very simple to integrate my code to this repo. It also works with functionDef. Would you like me to create a PR?

+class IdentifierTransformer(ast.NodeTransformer):
+    def __init__(self, identifiers=None):
+        self.identifiers = {} if identifiers is None else identifiers
+
+    def visit(self, node):
+        for field, value in iter_fields(node):
+            if isinstance(value, str):
+                value = self.identifiers.get(value, value)
+                setattr(node, field, value)
+        return super().visit(node)


def my_latexify(fn, identifiers=None):
    try:
        source = inspect.getsource(fn)
    # pylint: disable=broad-except
    except Exception:
        # Maybe running on console.
        source = dill.source.getsource(fn)

    source = textwrap.dedent(source)
    tree = ast.parse(source, mode="exec")
+    tree = fix_missing_locations(IdentifierTransformer(identifiers).visit(tree))

    return LatexifyVisitor().visit(tree)

from latexify_py.

@cccntu Ah sorry, I understand now that you meant to modify the tree with a transformer before converting the tree to LaTeX.

I think your example modifies every value that matches a specific entry in identifiers. E.g., if we give identifiers={"x": "y"} and the function has a literal "x", it is also converted.

I think we can follow the standard usage of NodeTransformer to constrain the application of rules. The following example converts only the identifiers appeared in the function:

import ast
import inspect

class NameReplacer(ast.NodeTransformer):
    def __init__(self, mapping: dict[str, str]):
        self._mapping = mapping

    def visit_Name(self, node: ast.Name) -> ast.Name:
        replaced = self._mapping.get(node.id, node.id)
        return ast.Name(id=replaced, ctx=node.ctx)

def f(x):
    return x

tree = ast.parse(inspect.getsource(f))
print(ast.dump(tree, indent=2))
converted = NameReplacer({"x": "y"}).visit(tree)
print(ast.dump(converted, indent=2))

This code outputs:

Module(
  body=[
    FunctionDef(
      name='f',
      args=arguments(
        posonlyargs=[],
        args=[
          arg(arg='x')],
        kwonlyargs=[],
        kw_defaults=[],
        defaults=[]),
      body=[
        Return(
          value=Name(id='x', ctx=Load()))],
      decorator_list=[])],
  type_ignores=[])
Module(
  body=[
    FunctionDef(
      name='f',
      args=arguments(
        posonlyargs=[],
        args=[
          arg(arg='x')],
        kwonlyargs=[],
        kw_defaults=[],
        defaults=[]),
      body=[
        Return(
          value=Name(id='y', ctx=Load()))],
      decorator_list=[])],
  type_ignores=[])

Note that this does not replace the function name and the argument list. To this end we also need to implement visit_FunctionDef appropriately (but this function is sensitive against the runtime version).

from latexify_py.

I think this is the complete version that supports both FunctionDef and Name (for Python >=3.8).

class NameReplacer(ast.NodeTransformer):
    def __init__(self, mapping: dict[str, str]):
        self._mapping = mapping

    def _replace_args(self, args: list[ast.arg]) -> list[ast.arg]:
        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]:
        return [self.visit(child) for child in children]

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
        return ast.FunctionDef(
            name=self._mapping.get(node.name, node.name),
            args=ast.arguments(
                posonlyargs=self._replace_args(node.args.posonlyargs),
                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),
            ),
            body=self._visit_children(node.body),
            decorator_list=self._visit_children(node.decorator_list),
        )

    def visit_Name(self, node: ast.Name) -> ast.Name:
        return ast.Name(
            id=self._mapping.get(node.id, node.id),
            ctx=node.ctx,
        )

from latexify_py.

Test:

a = b = c = 1

def d(fn):
    return fn

@d
def f(x=a, /, y=b, *, z=c):
    return x

tree = ast.parse(inspect.getsource(f))
print(ast.dump(tree, indent=2))
print()

identifiers = {x: x.upper() for x in "abcdfxyz"}
converted = NameReplacer(identifiers).visit(tree)
print(ast.dump(converted, indent=2))

All names are replaced appropriately:

Module(
  body=[
    FunctionDef(
      name='f',
      args=arguments(
        posonlyargs=[
          arg(arg='x')],
        args=[
          arg(arg='y')],
        kwonlyargs=[
          arg(arg='z')],
        kw_defaults=[
          Name(id='c', ctx=Load())],
        defaults=[
          Name(id='a', ctx=Load()),
          Name(id='b', ctx=Load())]),
      body=[
        Return(
          value=Name(id='x', ctx=Load()))],
      decorator_list=[
        Name(id='d', ctx=Load())])],
  type_ignores=[])

Module(
  body=[
    FunctionDef(
      name='F',
      args=arguments(
        posonlyargs=[
          arg(arg='X')],
        args=[
          arg(arg='Y')],
        kwonlyargs=[
          arg(arg='Z')],
        kw_defaults=[
          Name(id='C', ctx=Load())],
        defaults=[
          Name(id='A', ctx=Load()),
          Name(id='B', ctx=Load())]),
      body=[
        Return(
          value=Name(id='X', ctx=Load()))],
      decorator_list=[
        Name(id='D', ctx=Load())])],
  type_ignores=[])

from latexify_py.

Btw, I like the concept of applying conversion rules between ast.parse and LatexifyVisitor. We can introduce arbitrary number of conversion rules as follows:

tree = ast.parse(src)
tree = FooTransformer().visit(tree) if enable_foo else tree
tree = BarTransformer().visit(tree) if enable_bar else tree
...
latex = LatexifyVisitor(...).visit(tree)

This is really good outlook (though calculation cost increases), and we can also implement some existing rules as a transformer.

from latexify_py.

E.g., if we give identifiers={"x": "y"} and the function has a literal "x", it is also converted.

Nice catch. Yeah, writing out all the rules seems better.

Btw, I tried to see if there is a way to replace non-name stuff, like torch.sqrt -> sqrt. But Attribute is hard to replace, doing sub-tree matching feels too complicated.
Another easier option is to do str.replace after inspect.getsource, we can make this a separate option so we only do it when user ask for it.
Or we can annotation inside function.

def latex_replace(python_object, string):
    return python_object

@with_latex(function_name="\epsilon_{\theta}", identifiers={"alpha_t_minus_1": "{\alpha}_{t-1}"})
def my_long_function_name(alpha_t_minus_1):
     return latex_replace(torch.sqrt, "sqrt")(alpha_t_minus_1)

And we can do something to the latex_replace wrapped objects with another node transformer.

from latexify_py.