Mypy compatibility

.github/workflows/CI.yml

@@ -143,3 +143,29 @@ jobs:
         run: |
             pip install coveralls
             coveralls --finish
+
+  types:
+    name: Check types
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash -l {0}
+    strategy:
+      matrix:
+        python-version: ['3.10']
+
+    steps:
+      - uses: actions/checkout@v3
+      - name: "Set up Python"
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+      - name: "Install PySR and all dependencies"
+        run: |
+            python -m pip install --upgrade pip
+            pip install -r requirements.txt
+            pip install mypy jax jaxlib torch
+            python setup.py install
+      - name: "Run mypy"
+        run: mypy --install-types --non-interactive pysr

mypy.ini

@@ -0,0 +1,8 @@
+[mypy]
+warn_return_any = True
+
+[mypy-sklearn.*]
+ignore_missing_imports = True
+
+[mypy-julia.*]
+ignore_missing_imports = True

pysr/export_latex.py

@@ -1,5 +1,5 @@
 """Functions to help export PySR equations to LaTeX."""
-from typing import List
+from typing import List, Optional, Tuple
 
 import pandas as pd
 import sympy
@@ -19,14 +19,16 @@ class PreciseLatexPrinter(LatexPrinter):
         return super()._print_Float(reduced_float)
 
 
-def sympy2latex(expr, prec=3, full_prec=True, **settings):
+def sympy2latex(expr, prec=3, full_prec=True, **settings) -> str:
     """Convert sympy expression to LaTeX with custom precision."""
     settings["full_prec"] = full_prec
     printer = PreciseLatexPrinter(settings=settings, prec=prec)
     return printer.doprint(expr)
 
 
-def generate_table_environment(columns=["equation", "complexity", "loss"]):
+def generate_table_environment(
+    columns: List[str] = ["equation", "complexity", "loss"]
+) -> Tuple[str, str]:
     margins = "c" * len(columns)
     column_map = {
         "complexity": "Complexity",
@@ -58,12 +60,12 @@ def generate_table_environment(columns=["equation", "complexity", "loss"]):
 
 def sympy2latextable(
     equations: pd.DataFrame,
-    indices: List[int] = None,
+    indices: Optional[List[int]] = None,
     precision: int = 3,
-    columns=["equation", "complexity", "loss", "score"],
+    columns: List[str] = ["equation", "complexity", "loss", "score"],
     max_equation_length: int = 50,
     output_variable_name: str = "y",
-):
+) -> str:
     """Generate a booktabs-style LaTeX table for a single set of equations."""
     assert isinstance(equations, pd.DataFrame)
 
@@ -71,7 +73,7 @@ def sympy2latextable(
     latex_table_content = []
 
     if indices is None:
-        indices = range(len(equations))
+        indices = list(equations.index)
 
     for i in indices:
         latex_equation = sympy2latex(
@@ -126,11 +128,11 @@ def sympy2latextable(
 
 def sympy2multilatextable(
     equations: List[pd.DataFrame],
-    indices: List[List[int]] = None,
+    indices: Optional[List[List[int]]] = None,
     precision: int = 3,
-    columns=["equation", "complexity", "loss", "score"],
-    output_variable_names: str = None,
-):
+    columns: List[str] = ["equation", "complexity", "loss", "score"],
+    output_variable_names: Optional[List[str]] = None,
+) -> str:
     """Generate multiple latex tables for a list of equation sets."""
     # TODO: Let user specify custom output variable
 

pysr/export_sympy.py

@@ -51,14 +51,14 @@ sympy_mappings = {
 
 
 def create_sympy_symbols(
-    feature_names_in: Optional[List[str]] = None,
+    feature_names_in: List[str],
 ) -> List[sympy.Symbol]:
     return [sympy.Symbol(variable) for variable in feature_names_in]
 
 
 def pysr2sympy(
     equation: str, *, extra_sympy_mappings: Optional[Dict[str, Callable]] = None
-) -> sympy.Expr:
+):
     local_sympy_mappings = {
         **(extra_sympy_mappings if extra_sympy_mappings else {}),
         **sympy_mappings,

pysr/feature_selection.py

@@ -2,7 +2,7 @@
 import numpy as np
 
 
-def run_feature_selection(X, y, select_k_features, random_state=None) -> np.ndarray:
+def run_feature_selection(X, y, select_k_features, random_state=None):
     """
     Find most important features.
 

pysr/sr.py

@@ -11,6 +11,7 @@ from datetime import datetime
 from io import StringIO
 from multiprocessing import cpu_count
 from pathlib import Path
+from typing import List, Optional
 
 import numpy as np
 import pandas as pd
@@ -1781,10 +1782,10 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         y,
         Xresampled=None,
         weights=None,
-        variable_names=None,
-        X_units=None,
-        y_units=None,
-    ):
+        variable_names: Optional[List[str]] = None,
+        X_units: Optional[List[str]] = None,
+        y_units: Optional[List[str]] = None,
+    ) -> "PySRRegressor":
         """
         Search for equations to fit the dataset and store them in `self.equations_`.
 
@@ -2371,7 +2372,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
         return "\n".join(preamble_string + [table_string])
 
 
-def idx_model_selection(equations: pd.DataFrame, model_selection: str) -> int:
+def idx_model_selection(equations: pd.DataFrame, model_selection: str):
     """Select an expression and return its index."""
     if model_selection == "accuracy":
         chosen_idx = equations["loss"].idxmin()