fix google#31: add a switch for processing function names

latexify/core.py

@@ -27,8 +27,9 @@
 class LatexifyVisitor(node_visitor_base.NodeVisitorBase):
   """Latexify AST visitor."""
 
-  def __init__(self, math_symbol):
+  def __init__(self, math_symbol=False, raw_func_name=False):
     self.math_symbol = math_symbol
+    self.raw_func_name = raw_func_name  # True:do not treat underline as label of subscript(#31)
     super().__init__()
 
   def _parse_math_symbols(self, val: str) -> str:
@@ -52,6 +53,8 @@ def visit_FunctionDef(self, node, action):  # pylint: disable=invalid-name
     del action
 
     name_str = r'\mathrm{' + str(node.name) + '}'
+    if self.raw_func_name:
+      name_str = name_str.replace(r"_", r"\_")  # fix #31
     arg_strs = [
         self._parse_math_symbols(str(arg.arg)) for arg in node.args.args]
     body_str = self.visit(node.body[0])
@@ -225,26 +228,26 @@ def visit_If(self, node, action):  # pylint: disable=invalid-name
     return latex + r', & \mathrm{otherwise} \end{array} \right.'
 
 
-def get_latex(fn, math_symbol=True):
+def get_latex(fn, *args, **kwargs):
   try:
     source = inspect.getsource(fn)
   # pylint: disable=broad-except
   except Exception:
     # Maybe running on console.
     source = dill.source.getsource(fn)
 
-  return LatexifyVisitor(math_symbol=math_symbol).visit(ast.parse(source))
+  return LatexifyVisitor(*args, **kwargs).visit(ast.parse(source))
 
 
-def with_latex(*args, math_symbol=True):
+def with_latex(*args, **kwargs):
   """Translate a function with latex representation."""
 
   class _LatexifiedFunction:
     """Function with latex representation."""
 
     def __init__(self, fn):
       self._fn = fn
-      self._str = get_latex(fn, math_symbol=math_symbol)
+      self._str = get_latex(fn, *args,**kwargs)
 
     @property
     def __doc__(self):