test: fix mypy errors

MilesCranmer · Nov 28, 2024 · df2db4a · df2db4a
1 parent b3473f3
commit df2db4a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/pysr/expression_specs.py b/pysr/expression_specs.py
@@ -1,10 +1,11 @@
 import copy
+import sys
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, List, NewType, Optional
 
-try:
+if sys.version_info >= (3, 10):
     from typing import TypeAlias
-except ImportError:
+else:
     from typing_extensions import TypeAlias
 
 import numpy as np

diff --git a/pysr/sr.py b/pysr/sr.py
@@ -1071,7 +1071,7 @@ def from_file(
             assert unary_operators is None
             assert n_features_in is None
             with open(pkl_filename, "rb") as f:
-                model: "pysr.sr.PySRRegressor" = pkl.load(f)
+                model: "PySRRegressor" = pkl.load(f)
 
             # Update any parameters if necessary, such as
             # extra_sympy_mappings: