From 41568044e6e67cccbad056be63f0fd5db8a564da Mon Sep 17 00:00:00 2001
From: Jonathan Dekermanjian <Jonathan.Dekermanjian@gmail.com>
Date: Sat, 14 Sep 2024 09:49:15 -0600
Subject: [PATCH] added cleaner method using re to escape underscores, added
 cleaner test to assert underscores are escaped

---
 pymc/printing.py       | 22 ++++------------------
 tests/test_printing.py |  8 ++------
 2 files changed, 6 insertions(+), 24 deletions(-)

diff --git a/pymc/printing.py b/pymc/printing.py
index 2215c6f573..ef417f3799 100644
--- a/pymc/printing.py
+++ b/pymc/printing.py
@@ -13,6 +13,8 @@
 #   limitations under the License.
 
 
+import re
+
 from functools import partial
 
 from pytensor.compile import SharedVariable
@@ -301,22 +303,6 @@ def _default_repr_pretty(obj: TensorVariable | Model, p, cycle):
 
 def _format_underscore(variable: str) -> str:
     """
-    formats variables with underscores in its name by prefixing underscores by '\\'
-    ---
-    Params:
-        variable: The string representation of the variable in the model
+    Escapes all unescaped underscores in the variable name for LaTeX representation.
     """
-    if "_" not in variable:
-        return variable
-
-    inds = [i for i, ltr in enumerate(variable) if ltr == "_"]
-    var_len_original = len(variable)
-    var_len = None
-    for ind in inds:
-        if var_len:
-            if var_len != var_len_original:
-                ind = ind + (var_len - var_len_original)
-        if variable[ind - 1 : ind] != "\\":
-            variable = variable[:ind] + "\\" + variable[ind:]
-            var_len = len(variable)
-    return variable
+    return re.sub(r"(?<!\\)_", r"\\_", variable)
diff --git a/tests/test_printing.py b/tests/test_printing.py
index 9afc3a090d..832699c20d 100644
--- a/tests/test_printing.py
+++ b/tests/test_printing.py
@@ -11,7 +11,6 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-import re
 
 import numpy as np
 
@@ -335,8 +334,5 @@ def test_latex_escaped_underscore(self):
         """
         model = self.simple_model()
         model_str = model.str_repr(formatting="latex")
-        underscores = re.finditer(r"_", model_str)
-        for match in underscores:
-            if match:
-                start = match.span(0)[0] - 1
-                assert model_str[start : start + 1] == "\\"
+        assert "\\_" in model_str
+        assert "_" not in model_str.replace("\\_", "")