test: more typing info

pysr/julia_helpers.py

@@ -1,9 +1,10 @@
 """Functions for initializing the Julia environment and installing deps."""
 
-from typing import Any, Callable, cast
+from typing import Any, Callable, Union, cast
 
 import numpy as np
 from juliacall import convert as jl_convert  # type: ignore
+from numpy.typing import NDArray
 
 from .deprecated import init_julia, install
 from .julia_import import jl
@@ -26,7 +27,7 @@ def _escape_filename(filename):
     return str_repr
 
 
-def _load_cluster_manager(cluster_manager):
+def _load_cluster_manager(cluster_manager: str):
     jl.seval(f"using ClusterManagers: addprocs_{cluster_manager}")
     return jl.seval(f"addprocs_{cluster_manager}")
 
@@ -37,13 +38,13 @@ def jl_array(x):
     return jl_convert(jl.Array, x)
 
 
-def jl_serialize(obj):
+def jl_serialize(obj: Any) -> NDArray[np.uint8]:
     buf = jl.IOBuffer()
     Serialization.serialize(buf, obj)
     return np.array(jl.take_b(buf))
 
 
-def jl_deserialize(s):
+def jl_deserialize(s: Union[NDArray[np.uint8], None]):
     if s is None:
         return s
     buf = jl.IOBuffer()

pysr/sr.py

@@ -667,19 +667,19 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
     ```
     """
 
-    equations_: Optional[Union[pd.DataFrame, List[pd.DataFrame]]]
+    equations_: Union[pd.DataFrame, List[pd.DataFrame], None]
     n_features_in_: int
     feature_names_in_: ArrayLike[str]
     display_feature_names_in_: ArrayLike[str]
-    X_units_: Optional[ArrayLike[str]]
-    y_units_: Optional[Union[str, ArrayLike[str]]]
+    X_units_: Union[ArrayLike[str], None]
+    y_units_: Union[str, ArrayLike[str], None]
     nout_: int
-    selection_mask_: Optional[NDArray[np.bool_]]
+    selection_mask_: Union[NDArray[np.bool_], None]
     tempdir_: Path
     equation_file_: Union[str, Path]
-    julia_state_stream_: Optional[NDArray[np.uint8]]
-    julia_options_stream_: Optional[NDArray[np.uint8]]
-    equation_file_contents_: Optional[List[pd.DataFrame]]
+    julia_state_stream_: Union[NDArray[np.uint8], None]
+    julia_options_stream_: Union[NDArray[np.uint8], None]
+    equation_file_contents_: Union[List[pd.DataFrame], None]
     show_pickle_warnings_: bool
 
     def __init__(