diff --git a/slim_trees/__init__.py b/slim_trees/__init__.py index ea20a8d..ec1ac75 100644 --- a/slim_trees/__init__.py +++ b/slim_trees/__init__.py @@ -75,7 +75,7 @@ def dump_sklearn_compressed( """ from slim_trees.sklearn_tree import dump - dump_compressed(model, file, compression, dump) # type: ignore + dump_compressed(model, file, compression, dump) # type: ignore def dumps_sklearn_compressed( diff --git a/slim_trees/pickling.py b/slim_trees/pickling.py index dfa1ce1..3bd5d55 100644 --- a/slim_trees/pickling.py +++ b/slim_trees/pickling.py @@ -67,14 +67,15 @@ def _unpack_compression_args( return compression_method, _get_default_kwargs(compression_method) raise ValueError("File must be a path or compression must not be None.") + @overload def dump_compressed( obj: Any, file: BinaryIO, compression: Union[str, dict], dump_function: Optional[Callable] = None, -): - ... +): ... + @overload def dump_compressed( @@ -82,8 +83,8 @@ def dump_compressed( file: Union[str, pathlib.Path], compression: Optional[Union[str, dict]] = None, dump_function: Optional[Callable] = None, -): - ... +): ... + def dump_compressed( obj: Any, @@ -145,16 +146,16 @@ def load_compressed( file: BinaryIO, compression: Union[str, dict], unpickler_class: type = pickle.Unpickler, -): - ... +): ... + @overload def load_compressed( file: Union[str, pathlib.Path], compression: Optional[Union[str, dict]] = None, unpickler_class: type = pickle.Unpickler, -): - ... +): ... + def load_compressed( file: Union[str, pathlib.Path, BinaryIO],