Assert variable names are alphanumeric

Author: MilesCranmer

pysr/sr.py

@@ -181,6 +181,13 @@ def _check_assertions(
                 raise ValueError(
                     f"Variable name {var_name} is already a function name."
                 )
+            # Check if alphanumeric only:
+            if not re.match(r"^[a-zA-Z0-9_]+$", var_name):
+                raise ValueError(
+                    f"Invalid variable name {var_name}. "
+                    "Only alphanumeric characters, numbers, "
+                    "and underscores are allowed."
+                )
 
 
 def best(*args, **kwargs):  # pragma: no cover

test/test.py

@@ -572,8 +572,31 @@ def test_sympy_function_fails_as_variable(self):
         model = PySRRegressor()
         X = np.random.randn(100, 2)
         y = np.random.randn(100)
-        with self.assertRaises(ValueError):
+        with self.assertRaises(ValueError) as cm:
             model.fit(X, y, variable_names=["x1", "N"])
+        self.assertIn(
+            "Variable name",
+            str(cm.exception)
+        )
+
+    def test_bad_variable_names_fail(self):
+        model = PySRRegressor()
+        X = np.random.randn(100, 1)
+        y = np.random.randn(100)
+
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y, variable_names=["Tr(Tij)"])
+        self.assertIn(
+            "Invalid variable name",
+            str(cm.exception)
+        )
+
+        with self.assertRaises(ValueError) as cm:
+            model.fit(X, y, variable_names=["f{c}"])
+        self.assertIn(
+            "Invalid variable name",
+            str(cm.exception)
+        )
 
     def test_pickle_with_temp_equation_file(self):
         """If we have a temporary equation file, unpickle the estimator."""