ENH: add escape math mode with escape=latex-math

natmokval · natmokval · commit 4f9e8e189cab · 2023-01-28T22:18:34.000+01:00
diff --git a/pandas/io/formats/style_render.py b/pandas/io/formats/style_render.py
@@ -989,6 +989,8 @@ def format(
             Use 'latex' to replace the characters ``&``, ``%``, ``$``, ``#``, ``_``,
             ``{``, ``}``, ``~``, ``^``, and ``\`` in the cell display string with
             LaTeX-safe sequences.
+            Use 'latex-math' to replace the characters the same way as in 'latex' mode,
+            except for math substrings, which start and end with ``$``.
             Escaping is done before ``formatter``.
 
             .. versionadded:: 1.3.0
@@ -1105,16 +1107,28 @@ def format(
         <td .. >NA</td>
         ...
 
-        Using a ``formatter`` with LaTeX ``escape``.
+        Using a ``formatter`` with ``escape`` in 'latex' mode.
 
-        >>> df = pd.DataFrame([["123"], ["~ ^"], ["$%#"]])
+        >>> df = pd.DataFrame([["123"], ["~ ^"], ["%#"]])
         >>> df.style.format("\\textbf{{{}}}", escape="latex").to_latex()
         ...  # doctest: +SKIP
         \begin{tabular}{ll}
-        {} & {0} \\
+         & 0 \\
         0 & \textbf{123} \\
         1 & \textbf{\textasciitilde \space \textasciicircum } \\
-        2 & \textbf{\$\%\#} \\
+        2 & \textbf{\%\#} \\
+        \end{tabular}
+
+        Using ``escape`` in 'latex-math' mode.
+
+        >>> df = pd.DataFrame([[r"$\sum_{i=1}^{10} a_i$ a~b $\alpha \
+        ...     = \frac{\beta}{\zeta^2}$"], ["%#^ $ \$x^2 $"]])
+        >>> df.style.format(escape="latex-math").to_latex()
+        ...  # doctest: +SKIP
+        \begin{tabular}{ll}
+         & 0 \\
+        0 & $\sum_{i=1}^{10} a_i$ a\textasciitilde b $\alpha = \frac{\beta}{\zeta^2}$ \\
+        1 & \%\#\textasciicircum \space $ \$x^2 $ \\
         \end{tabular}
 
         Pandas defines a `number-format` pseudo CSS attribute instead of the `.format`
@@ -1743,9 +1757,12 @@ def _str_escape(x, escape):
             return escape_html(x)
         elif escape == "latex":
             return _escape_latex(x)
+        elif escape == "latex-math":
+            return _escape_latex_math(x)
         else:
             raise ValueError(
-                f"`escape` only permitted in {{'html', 'latex'}}, got {escape}"
+                f"`escape` only permitted in {{'html', 'latex', 'latex-math'}}, \
+got {escape}"
             )
     return x
 
@@ -2344,3 +2361,36 @@ def _escape_latex(s):
         .replace("^", "\\textasciicircum ")
         .replace("ab2§=§8yz", "\\textbackslash ")
     )
+
+
+def _escape_latex_math(s):
+    r"""
+    All characters between two characters ``$`` are preserved.
+
+    The substrings in LaTeX math mode, which start with the character ``$``
+    and end with ``$``, are preserved without escaping. Otherwise
+    regular LaTeX escaping applies. See ``_escape_latex()``.
+
+    Parameters
+    ----------
+    s : str
+        Input to be escaped
+
+    Return
+    ------
+    str :
+        Escaped string
+    """
+    s = s.replace(r"\$", r"ab2§=§8yz")
+    pattern = re.compile(r"\$.*?\$")
+    pos = 0
+    ps = pattern.search(s, pos)
+    res = []
+    while ps:
+        res.append(_escape_latex(s[pos : ps.span()[0]]))
+        res.append(ps.group())
+        pos = ps.span()[1]
+        ps = pattern.search(s, pos)
+
+    res.append(_escape_latex(s[pos : len(s)]))
+    return "".join(res).replace(r"ab2§=§8yz", r"\$")
diff --git a/pandas/tests/io/formats/style/test_format.py b/pandas/tests/io/formats/style/test_format.py
@@ -359,7 +359,7 @@ def test_format_decimal(formatter, thousands, precision, func, col):
 
 
 def test_str_escape_error():
-    msg = "`escape` only permitted in {'html', 'latex'}, got "
+    msg = "`escape` only permitted in {'html', 'latex', 'latex-math'}, got "
     with pytest.raises(ValueError, match=msg):
         _str_escape("text", "bad_escape")
 
