Try to fix nb sanitizer

pysr/_cli/main.py

@@ -1,3 +1,4 @@
+import unittest
 import warnings
 
 import click
@@ -55,19 +56,27 @@ def _tests(tests):
 
     Choose from main, jax, torch, cli, dev, and startup. You can give multiple tests, separated by commas.
     """
+    test_cases = []
     for test in tests.split(","):
         if test == "main":
-            runtests()
+            test_cases.extend(runtests(just_tests=True))
         elif test == "jax":
-            runtests_jax()
+            test_cases.extend(runtests_jax(just_tests=True))
         elif test == "torch":
-            runtests_torch()
+            test_cases.extend(runtests_torch(just_tests=True))
         elif test == "cli":
             runtests_cli = get_runtests_cli()
-            runtests_cli()
+            test_cases.extend(runtests_cli(just_tests=True))
         elif test == "dev":
-            runtests_dev()
+            test_cases.extend(runtests_dev(just_tests=True))
         elif test == "startup":
-            runtests_startup()
+            test_cases.extend(runtests_startup(just_tests=True))
         else:
             warnings.warn(f"Invalid test {test}. Skipping.")
+
+    loader = unittest.TestLoader()
+    suite = unittest.TestSuite()
+    for test_case in test_cases:
+        suite.addTests(loader.loadTestsFromTestCase(test_case))
+    runner = unittest.TextTestRunner()
+    return runner.run(suite)

pysr/test/test.py

@@ -1127,10 +1127,8 @@ class TestDimensionalConstraints(unittest.TestCase):
 # TODO: Determine desired behavior if second .fit() call does not have units
 
 
-def runtests():
+def runtests(just_tests=False):
     """Run all tests in test.py."""
-    suite = unittest.TestSuite()
-    loader = unittest.TestLoader()
     test_cases = [
         TestPipeline,
         TestBest,
@@ -1139,8 +1137,11 @@ def runtests():
         TestLaTeXTable,
         TestDimensionalConstraints,
     ]
+    if just_tests:
+        return test_cases
+    suite = unittest.TestSuite()
+    loader = unittest.TestLoader()
     for test_case in test_cases:
-        tests = loader.loadTestsFromTestCase(test_case)
-        suite.addTests(tests)
+        suite.addTests(loader.loadTestsFromTestCase(test_case))
     runner = unittest.TextTestRunner()
     return runner.run(suite)

pysr/test/test_cli.py

@@ -68,11 +68,15 @@ def get_runtests():
             self.assertEqual(result.output.strip(), expected.strip())
             self.assertEqual(result.exit_code, 0)
 
-    def runtests():
+    def runtests(just_tests=False):
         """Run all tests in cliTest.py."""
+        tests = [TestCli]
+        if just_tests:
+            return tests
         loader = unittest.TestLoader()
         suite = unittest.TestSuite()
-        suite.addTests(loader.loadTestsFromTestCase(TestCli))
+        for test in tests:
+            suite.addTests(loader.loadTestsFromTestCase(test))
         runner = unittest.TextTestRunner()
         return runner.run(suite)
 

pysr/test/test_dev.py

@@ -47,9 +47,13 @@ class TestDev(unittest.TestCase):
         self.assertEqual(test_result.stdout.decode("utf-8").strip(), "2.3")
 
 
-def runtests():
+def runtests(just_tests=False):
+    tests = [TestDev]
+    if just_tests:
+        return tests
     suite = unittest.TestSuite()
     loader = unittest.TestLoader()
-    suite.addTests(loader.loadTestsFromTestCase(TestDev))
+    for test in tests:
+        suite.addTests(loader.loadTestsFromTestCase(test))
     runner = unittest.TextTestRunner()
     return runner.run(suite)

pysr/test/test_jax.py

@@ -121,10 +121,14 @@ class TestJAX(unittest.TestCase):
         np.testing.assert_almost_equal(y.values, jax_output, decimal=3)
 
 
-def runtests():
+def runtests(just_tests=False):
     """Run all tests in test_jax.py."""
+    tests = [TestJAX]
+    if just_tests:
+        return tests
     loader = unittest.TestLoader()
     suite = unittest.TestSuite()
-    suite.addTests(loader.loadTestsFromTestCase(TestJAX))
+    for test in tests:
+        suite.addTests(loader.loadTestsFromTestCase(test))
     runner = unittest.TextTestRunner()
     return runner.run(suite)

pysr/test/test_startup.py

@@ -143,9 +143,13 @@ class TestStartup(unittest.TestCase):
         self.assertEqual(result.returncode, 0)
 
 
-def runtests():
+def runtests(just_tests=False):
+    tests = [TestStartup]
+    if just_tests:
+        return tests
     suite = unittest.TestSuite()
     loader = unittest.TestLoader()
-    suite.addTests(loader.loadTestsFromTestCase(TestStartup))
+    for test in tests:
+        suite.addTests(loader.loadTestsFromTestCase(test))
     runner = unittest.TextTestRunner()
     return runner.run(suite)

pysr/test/test_torch.py

@@ -184,10 +184,14 @@ class TestTorch(unittest.TestCase):
         np.testing.assert_almost_equal(y.values, torch_output, decimal=3)
 
 
-def runtests():
+def runtests(just_tests=False):
     """Run all tests in test_torch.py."""
+    tests = [TestTorch]
+    if just_tests:
+        return tests
     loader = unittest.TestLoader()
     suite = unittest.TestSuite()
-    suite.addTests(loader.loadTestsFromTestCase(TestTorch))
+    for test in tests:
+        suite.addTests(loader.loadTestsFromTestCase(test))
     runner = unittest.TextTestRunner()
     return runner.run(suite)