Add built-in fn enum (#130)

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

Author: ZibingZhang

src/latexify/constants.py

@@ -2,39 +2,79 @@
 
 from __future__ import annotations
 
+import enum
+
 PREFIXES = ["math", "numpy", "np"]
 
-BUILTIN_FUNCS: dict[str, tuple[str, str]] = {
-    "abs": (r"\left|{", r"}\right|"),
-    "acos": (r"\arccos{\left({", r"}\right)}"),
-    "acosh": (r"\mathrm{arccosh}{\left({", r"}\right)}"),
-    "arccos": (r"\arccos{\left({", r"}\right)}"),
-    "arccosh": (r"\mathrm{arccosh}{\left({", r"}\right)}"),
-    "arcsin": (r"\arcsin{\left({", r"}\right)}"),
-    "arcsinh": (r"\mathrm{arcsinh}{\left({", r"}\right)}"),
-    "arctan": (r"\arctan{\left({", r"}\right)}"),
-    "arctanh": (r"\mathrm{arctanh}{\left({", r"}\right)}"),
-    "asin": (r"\arcsin{\left({", r"}\right)}"),
-    "asinh": (r"\mathrm{arcsinh}{\left({", r"}\right)}"),
-    "atan": (r"\arctan{\left({", r"}\right)}"),
-    "atanh": (r"\mathrm{arctanh}{\left({", r"}\right)}"),
-    "ceil": (r"\left\lceil{", r"}\right\rceil"),
-    "cos": (r"\cos{\left({", r"}\right)}"),
-    "cosh": (r"\cosh{\left({", r"}\right)}"),
-    "exp": (r"\exp{\left({", r"}\right)}"),
-    "fabs": (r"\left|{", r"}\right|"),
-    "factorial": (r"\left({", r"}\right)!"),
-    "floor": (r"\left\lfloor{", r"}\right\rfloor"),
-    "fsum": (r"\sum\left({", r"}\right)"),
-    "gamma": (r"\Gamma\left({", r"}\right)"),
-    "log": (r"\log{\left({", r"}\right)}"),
-    "log10": (r"\log_{10}{\left({", r"}\right)}"),
-    "log2": (r"\log_{2}{\left({", r"}\right)}"),
-    "prod": (r"\prod \left({", r"}\right)"),
-    "sin": (r"\sin{\left({", r"}\right)}"),
-    "sinh": (r"\sinh{\left({", r"}\right)}"),
-    "sqrt": (r"\sqrt{", "}"),
-    "tan": (r"\tan{\left({", r"}\right)}"),
-    "tanh": (r"\tanh{\left({", r"}\right)}"),
-    "sum": (r"\sum \left({", r"}\right)"),
+
+class BuiltinFnName(str, enum.Enum):
+    """Built-in function name."""
+
+    ABS = "abs"
+    ACOS = "acos"
+    ACOSH = "acosh"
+    ARCCOS = "arccos"
+    ARCCOSH = "arcosh"
+    ARCSIN = "arcsin"
+    ARCSINH = "arcsihn"
+    ARCTAN = "arctan"
+    ARCTANH = "arctanh"
+    ASIN = "asin"
+    ASINH = "asinh"
+    ATAN = "atan"
+    ATANH = "atanh"
+    CEIL = "ceil"
+    COS = "cos"
+    COSH = "cosh"
+    EXP = "exp"
+    FABS = "fabs"
+    FACTORIAL = "factorial"
+    FLOOR = "floor"
+    FSUM = "fsum"
+    GAMMA = "gamma"
+    LOG = "log"
+    LOG10 = "log10"
+    LOG2 = "log2"
+    PROD = "prod"
+    SIN = "sin"
+    SINH = "sinh"
+    SQRT = "sqrt"
+    TAN = "tan"
+    TANH = "tanh"
+    SUM = "sum"
+
+
+BUILTIN_FUNCS: dict[BuiltinFnName, tuple[str, str]] = {
+    BuiltinFnName.ABS: (r"\left|{", r"}\right|"),
+    BuiltinFnName.ACOS: (r"\arccos{\left({", r"}\right)}"),
+    BuiltinFnName.ACOSH: (r"\mathrm{arccosh}{\left({", r"}\right)}"),
+    BuiltinFnName.ARCCOS: (r"\arccos{\left({", r"}\right)}"),
+    BuiltinFnName.ARCCOSH: (r"\mathrm{arccosh}{\left({", r"}\right)}"),
+    BuiltinFnName.ARCSIN: (r"\arcsin{\left({", r"}\right)}"),
+    BuiltinFnName.ARCSINH: (r"\mathrm{arcsinh}{\left({", r"}\right)}"),
+    BuiltinFnName.ARCTAN: (r"\arctan{\left({", r"}\right)}"),
+    BuiltinFnName.ARCTANH: (r"\mathrm{arctanh}{\left({", r"}\right)}"),
+    BuiltinFnName.ASIN: (r"\arcsin{\left({", r"}\right)}"),
+    BuiltinFnName.ASINH: (r"\mathrm{arcsinh}{\left({", r"}\right)}"),
+    BuiltinFnName.ATAN: (r"\arctan{\left({", r"}\right)}"),
+    BuiltinFnName.ATANH: (r"\mathrm{arctanh}{\left({", r"}\right)}"),
+    BuiltinFnName.CEIL: (r"\left\lceil{", r"}\right\rceil"),
+    BuiltinFnName.COS: (r"\cos{\left({", r"}\right)}"),
+    BuiltinFnName.COSH: (r"\cosh{\left({", r"}\right)}"),
+    BuiltinFnName.EXP: (r"\exp{\left({", r"}\right)}"),
+    BuiltinFnName.FABS: (r"\left|{", r"}\right|"),
+    BuiltinFnName.FACTORIAL: (r"\left({", r"}\right)!"),
+    BuiltinFnName.FLOOR: (r"\left\lfloor{", r"}\right\rfloor"),
+    BuiltinFnName.FSUM: (r"\sum\left({", r"}\right)"),
+    BuiltinFnName.GAMMA: (r"\Gamma\left({", r"}\right)"),
+    BuiltinFnName.LOG: (r"\log{\left({", r"}\right)}"),
+    BuiltinFnName.LOG10: (r"\log_{10}{\left({", r"}\right)}"),
+    BuiltinFnName.LOG2: (r"\log_{2}{\left({", r"}\right)}"),
+    BuiltinFnName.PROD: (r"\prod \left({", r"}\right)"),
+    BuiltinFnName.SIN: (r"\sin{\left({", r"}\right)}"),
+    BuiltinFnName.SINH: (r"\sinh{\left({", r"}\right)}"),
+    BuiltinFnName.SQRT: (r"\sqrt{", "}"),
+    BuiltinFnName.TAN: (r"\tan{\left({", r"}\right)}"),
+    BuiltinFnName.TANH: (r"\tanh{\left({", r"}\right)}"),
+    BuiltinFnName.SUM: (r"\sum \left({", r"}\right)"),
 }

src/latexify/transformers/function_expander.py

@@ -4,7 +4,7 @@
 import functools
 from collections.abc import Callable
 
-from latexify import ast_utils
+from latexify import ast_utils, constants
 
 
 # TODO(ZibingZhang): handle recursive function expansions
@@ -50,7 +50,7 @@ def _hypot_expander(function_expander: FunctionExpander, node: ast.Call) -> ast.
 
     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()),
+        func=ast.Name(id=constants.BuiltinFnName.SQRT.value, ctx=ast.Load()),
         args=[args_reduced],
     )