MilesCranmer commited on
Commit
4bcc280
1 Parent(s): 018c3a1

refactor: create `ArrayLike` type for type checking

Browse files
Files changed (1) hide show
  1. pysr/utils.py +8 -2
pysr/utils.py CHANGED
@@ -1,10 +1,16 @@
1
  import os
2
  import re
 
 
3
 
4
- from sklearn.utils.validation import _check_feature_names_in
 
5
 
 
 
6
 
7
- def _csv_filename_to_pkl_filename(csv_filename: str) -> str:
 
8
  if os.path.splitext(csv_filename)[1] == ".pkl":
9
  return csv_filename
10
 
 
1
  import os
2
  import re
3
+ from pathlib import Path
4
+ from typing import Any, List, TypeVar, Union
5
 
6
+ from numpy.typing import NDArray
7
+ from sklearn.utils.validation import _check_feature_names_in # type: ignore
8
 
9
+ T = TypeVar("T", bound=Any)
10
+ ArrayLike = Union[NDArray[T], List[T]]
11
 
12
+
13
+ def _csv_filename_to_pkl_filename(csv_filename: Union[str, Path]) -> Union[str, Path]:
14
  if os.path.splitext(csv_filename)[1] == ".pkl":
15
  return csv_filename
16