Merge pull request #834 from MilesCranmer/fix-sklearn-test

test: update deprecated sklearn test syntax

pysr/test/test_main.py

@@ -1,3 +1,4 @@
+import functools
 import importlib
 import os
 import pickle as pkl
@@ -12,7 +13,13 @@
 import numpy as np
 import pandas as pd
 import sympy  # type: ignore
-from sklearn.utils.estimator_checks import check_estimator
+
+try:
+    from sklearn.utils.estimator_checks import estimator_checks_generator
+except ImportError:
+    from sklearn.utils.estimator_checks import check_estimator
+
+    estimator_checks_generator = functools.partial(check_estimator, generate_only=True)
 
 from pysr import (
     ParametricExpressionSpec,
@@ -930,9 +937,8 @@ def test_scikit_learn_compatibility(self):
             temp_equation_file=True,
         )  # Return early.
 
-        check_generator = check_estimator(model, generate_only=True)
         exception_messages = []
-        for _, check in check_generator:
+        for _, check in estimator_checks_generator(model):
             if check.func.__name__ in {
                 # We can use complex data, so avoid this check
                 "check_complex_data",