Merge pull request #670 from MilesCranmer/issue666

fix: `from pysr import *`
MilesCranmer · Jul 15, 2024 · b658d24 · b658d24
2 parents db44938 + e84bed4
commit b658d24
Show file tree

Hide file tree

Showing 9 changed files with 21 additions and 13 deletions.
diff --git a/pysr/__init__.py b/pysr/__init__.py
@@ -18,7 +18,6 @@
     "sklearn_monkeypatch",
     "sympy2jax",
     "sympy2torch",
-    "Problem",
     "install",
     "PySRRegressor",
     "best",

diff --git a/pysr/export_jax.py b/pysr/export_jax.py
@@ -1,5 +1,5 @@
 import numpy as np  # noqa: F401
-import sympy
+import sympy  # type: ignore
 
 # Special since need to reduce arguments.
 MUL = 0

diff --git a/pysr/export_latex.py b/pysr/export_latex.py
@@ -3,8 +3,8 @@
 from typing import List, Optional, Tuple
 
 import pandas as pd
-import sympy
-from sympy.printing.latex import LatexPrinter
+import sympy  # type: ignore
+from sympy.printing.latex import LatexPrinter  # type: ignore
 
 
 class PreciseLatexPrinter(LatexPrinter):

diff --git a/pysr/export_numpy.py b/pysr/export_numpy.py
@@ -6,7 +6,7 @@
 import numpy as np
 import pandas as pd
 from numpy.typing import NDArray
-from sympy import Expr, Symbol, lambdify
+from sympy import Expr, Symbol, lambdify  # type: ignore
 
 
 def sympy2numpy(eqn, sympy_symbols, *, selection=None):

diff --git a/pysr/export_sympy.py b/pysr/export_sympy.py
@@ -2,7 +2,7 @@
 
 from typing import Callable, Dict, List, Optional
 
-import sympy
+import sympy  # type: ignore
 from sympy import sympify
 
 from .utils import ArrayLike

diff --git a/pysr/export_torch.py b/pysr/export_torch.py
@@ -4,7 +4,7 @@
 import functools as ft
 
 import numpy as np  # noqa: F401
-import sympy
+import sympy  # type: ignore
 
 
 def _reduce(fn):

diff --git a/pysr/test/test.py b/pysr/test/test.py
@@ -1,3 +1,4 @@
+import importlib
 import os
 import pickle as pkl
 import tempfile
@@ -8,7 +9,7 @@
 
 import numpy as np
 import pandas as pd
-import sympy
+import sympy  # type: ignore
 from sklearn.utils.estimator_checks import check_estimator
 
 from pysr import PySRRegressor, install, jl
@@ -892,7 +893,7 @@ def test_suggest_keywords(self):
 
         # More complex, and with error
         with self.assertRaises(TypeError) as cm:
-            model = PySRRegressor(ncyclesperiterationn=5)
+            PySRRegressor(ncyclesperiterationn=5)
 
         self.assertIn(
             "`ncyclesperiterationn` is not a valid keyword", str(cm.exception)
@@ -903,10 +904,18 @@ def test_suggest_keywords(self):
 
         # Farther matches (this might need to be changed)
         with self.assertRaises(TypeError) as cm:
-            model = PySRRegressor(operators=["+", "-"])
+            PySRRegressor(operators=["+", "-"])
 
         self.assertIn("`unary_operators`, `binary_operators`", str(cm.exception))
 
+    def test_issue_666(self):
+        # Try the equivalent of `from pysr import *`
+        pysr_module = importlib.import_module("pysr")
+        names_to_import = pysr_module.__all__
+
+        for name in names_to_import:
+            getattr(pysr_module, name)
+
 
 TRUE_PREAMBLE = "\n".join(
     [

diff --git a/pysr/test/test_jax.py b/pysr/test/test_jax.py
@@ -3,7 +3,7 @@
 
 import numpy as np
 import pandas as pd
-import sympy
+import sympy  # type: ignore
 
 import pysr
 from pysr import PySRRegressor, sympy2jax
@@ -102,7 +102,7 @@ def test_avoid_simplification(self):
         )
 
     def test_issue_656(self):
-        import sympy
+        import sympy  # type: ignore
 
         E_plus_x1 = sympy.exp(1) + sympy.symbols("x1")
         f, params = pysr.export_jax.sympy2jax(E_plus_x1, [sympy.symbols("x1")])

diff --git a/pysr/test/test_torch.py b/pysr/test/test_torch.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 import pandas as pd
-import sympy
+import sympy  # type: ignore
 
 import pysr
 from pysr import PySRRegressor, sympy2torch