fix: respect tempdir for temporary equation files

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

Author: MilesCranmerBot

pysr/sr.py

@@ -1590,15 +1590,16 @@ def _setup_equation_file(self):
             if self.output_directory is not None:
                 assert self.output_directory_ == self.output_directory
         else:
-            self.output_directory_ = (
-                tempfile.mkdtemp()
-                if self.temp_equation_file
-                else (
+            if self.temp_equation_file:
+                if self.tempdir is not None:
+                    Path(self.tempdir).mkdir(parents=True, exist_ok=True)
+                self.output_directory_ = tempfile.mkdtemp(dir=self.tempdir)
+            else:
+                self.output_directory_ = (
                     "outputs"
                     if self.output_directory is None
                     else self.output_directory
                 )
-            )
             self.run_id_ = (
                 cast(str, SymbolicRegression.SearchUtilsModule.generate_run_id())
                 if self.run_id is None

pysr/test/test_main.py

@@ -77,6 +77,13 @@ def setUp(self):
         self.rstate = np.random.RandomState(0)
         self.X = self.rstate.randn(100, 5)
 
+    def test_temp_equation_file_respects_tempdir(self):
+        with tempfile.TemporaryDirectory() as d:
+            tempdir = Path(d) / "pysr-temp"
+            model = PySRRegressor(temp_equation_file=True, tempdir=str(tempdir), run_id="t")
+            model._setup_equation_file()
+            self.assertEqual(Path(model.output_directory_).parent, tempdir)
+
     def test_linear_relation(self):
         y = self.X[:, 0]
         model = PySRRegressor(