Overhaul type hints with beartype and PEP 585 compliance

pysr/__init__.py

@@ -1,3 +1,13 @@
+try:
+ from beartype import BeartypeConf
+ from beartype.claw import beartype_this_package
+
+ beartype_this_package(conf=BeartypeConf(violation_type=UserWarning))
+except ImportError:
+ import warnings
+
+ warnings.warn("beartype not installed, skipping runtime type checks.")
+
 # This must be imported as early as possible to prevent
 # library linking issues caused by numpy/pytorch/etc. importing
 # old libraries:

pysr/export_latex.py

@@ -1,5 +1,5 @@
 """Functions to help export PySR equations to LaTeX."""
-from typing import List, Optional, Tuple
+from __future__ import annotations
 
 import pandas as pd
 import sympy
@@ -27,8 +27,8 @@ def sympy2latex(expr, prec=3, full_prec=True, **settings) -> str:
 
 
 def generate_table_environment(
- columns: List[str] = ["equation", "complexity", "loss"]
-) -> Tuple[str, str]:
+ columns: list[str] = ["equation", "complexity", "loss"]
+) -> tuple[str, str]:
  margins = "c" * len(columns)
  column_map = {
  "complexity": "Complexity",
@@ -60,9 +60,9 @@ def generate_table_environment(
 
 def sympy2latextable(
  equations: pd.DataFrame,
- indices: Optional[List[int]] = None,
+ indices: list[int] | None = None,
  precision: int = 3,
- columns: List[str] = ["equation", "complexity", "loss", "score"],
+ columns: list[str] = ["equation", "complexity", "loss", "score"],
  max_equation_length: int = 50,
  output_variable_name: str = "y",
 ) -> str:
@@ -127,11 +127,11 @@ def sympy2latextable(
 
 
 def sympy2multilatextable(
- equations: List[pd.DataFrame],
- indices: Optional[List[List[int]]] = None,
+ equations: list[pd.DataFrame],
+ indices: list[list[int]] | None = None,
  precision: int = 3,
- columns: List[str] = ["equation", "complexity", "loss", "score"],
- output_variable_names: Optional[List[str]] = None,
+ columns: list[str] = ["equation", "complexity", "loss", "score"],
+ output_variable_names: list[str] | None = None,
 ) -> str:
  """Generate multiple latex tables for a list of equation sets."""
  # TODO: Let user specify custom output variable

pysr/export_sympy.py

@@ -1,5 +1,7 @@
 """Define utilities to export to sympy"""
-from typing import Callable, Dict, List, Optional
+from __future__ import annotations
+
+from collections.abc import Callable
 
 import sympy
 from sympy import sympify
@@ -58,13 +60,13 @@
 
 
 def create_sympy_symbols(
- feature_names_in: List[str],
-) -> List[sympy.Symbol]:
+ 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
+ equation: str, *, extra_sympy_mappings: dict[str, Callable] | None = None
 ):
  local_sympy_mappings = {
  **(extra_sympy_mappings if extra_sympy_mappings else {}),

pysr/julia_helpers.py

@@ -15,11 +15,9 @@
 jl.seval("using SymbolicRegression: plus, sub, mult, div, pow")
 
 
-def _escape_filename(filename):
+def _escape_filename(filename: str) -> str:
  """Turn a path into a string with correctly escaped backslashes."""
- str_repr = str(filename)
- str_repr = str_repr.replace("\\", "\\\\")
- return str_repr
+ return filename.replace("\\", "\\\\")
 
 
 def _load_cluster_manager(cluster_manager):