refactor: typings for sympy export

MilesCranmer · Apr 28, 2024 · 583beaf · 583beaf
1 parent 7113eed
commit 583beaf
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/pysr/export_sympy.py b/pysr/export_sympy.py
@@ -5,6 +5,8 @@
 import sympy
 from sympy import sympify
 
+from .utils import ArrayLike
+
 sympy_mappings = {
     "div": lambda x, y: x / y,
     "mult": lambda x, y: x * y,
@@ -60,21 +62,21 @@
 
 
 def create_sympy_symbols_map(
-    feature_names_in: List[str],
+    feature_names_in: ArrayLike[str],
 ) -> Dict[str, sympy.Symbol]:
     return {variable: sympy.Symbol(variable) for variable in feature_names_in}
 
 
 def create_sympy_symbols(
-    feature_names_in: List[str],
+    feature_names_in: ArrayLike[str],
 ) -> List[sympy.Symbol]:
     return [sympy.Symbol(variable) for variable in feature_names_in]
 
 
 def pysr2sympy(
     equation: str,
     *,
-    feature_names_in: Optional[List[str]] = None,
+    feature_names_in: Optional[ArrayLike[str]] = None,
     extra_sympy_mappings: Optional[Dict[str, Callable]] = None,
 ):
     if feature_names_in is None: