diff --git a/src/pystow/api.py b/src/pystow/api.py index 81b4017..1372098 100644 --- a/src/pystow/api.py +++ b/src/pystow/api.py @@ -3,6 +3,8 @@ """API functions for PyStow.""" import bz2 +import io +import lzma import sqlite3 from contextlib import contextmanager from io import BytesIO, StringIO @@ -19,7 +21,7 @@ overload, ) -from .constants import JSON, BytesOpener, Opener, Provider +from .constants import JSON, BytesOpener, Provider from .impl import Module if TYPE_CHECKING: @@ -503,6 +505,34 @@ def ensure_open_zip( yield yv +@overload +@contextmanager +def ensure_open_lzma( + key: str, + *subkeys: str, + url: str, + name: Optional[str], + force: bool, + download_kwargs: Optional[Mapping[str, Any]], + mode: Literal["r", "w", "rt", "wt"] = "rt", + open_kwargs: Optional[Mapping[str, Any]], +) -> Generator[io.TextIOWrapper[lzma.LZMAFile], None, None]: ... + + +@overload +@contextmanager +def ensure_open_lzma( + key: str, + *subkeys: str, + url: str, + name: Optional[str], + force: bool, + download_kwargs: Optional[Mapping[str, Any]], + mode: Literal["rb", "wb"] = ..., + open_kwargs: Optional[Mapping[str, Any]], +) -> Generator[lzma.LZMAFile, None, None]: ... + + @contextmanager def ensure_open_lzma( key: str, @@ -511,9 +541,9 @@ def ensure_open_lzma( name: Optional[str] = None, force: bool = False, download_kwargs: Optional[Mapping[str, Any]] = None, - mode: str = "r", + mode: Literal["r", "rb", "w", "wb", "rt", "wt"] = "rt", open_kwargs: Optional[Mapping[str, Any]] = None, -) -> Opener: +) -> Generator[Union[lzma.LZMAFile, io.TextIOWrapper[lzma.LZMAFile]], None, None]: """Ensure a LZMA-compressed file is downloaded and open a file inside it. :param key: diff --git a/src/pystow/impl.py b/src/pystow/impl.py index 94ffd29..255b9a0 100644 --- a/src/pystow/impl.py +++ b/src/pystow/impl.py @@ -4,6 +4,7 @@ import bz2 import gzip +import io import json import logging import lzma @@ -27,7 +28,7 @@ ) from . import utils -from .constants import JSON, BytesOpener, Opener, Provider +from .constants import JSON, BytesOpener, Provider from .utils import ( base_from_gzip_name, download_from_google, @@ -485,6 +486,32 @@ def open_gz( with gzip.open(path, **open_kwargs) as file: yield file + @overload + @contextmanager + def ensure_open_lzma( + self, + *subkeys: str, + url: str, + name: Optional[str], + force: bool, + download_kwargs: Optional[Mapping[str, Any]], + mode: Literal["r", "w", "rt", "wt"] = "rt", + open_kwargs: Optional[Mapping[str, Any]], + ) -> Generator[io.TextIOWrapper[lzma.LZMAFile], None, None]: ... + + @overload + @contextmanager + def ensure_open_lzma( + self, + *subkeys: str, + url: str, + name: Optional[str], + force: bool, + download_kwargs: Optional[Mapping[str, Any]], + mode: Literal["rb", "wb"] = ..., + open_kwargs: Optional[Mapping[str, Any]], + ) -> Generator[lzma.LZMAFile, None, None]: ... + @contextmanager def ensure_open_lzma( self, @@ -493,9 +520,9 @@ def ensure_open_lzma( name: Optional[str] = None, force: bool = False, download_kwargs: Optional[Mapping[str, Any]] = None, - mode: str = "rt", + mode: Literal["r", "rb", "w", "wb", "rt", "wt"] = "rt", open_kwargs: Optional[Mapping[str, Any]] = None, - ) -> Opener: + ) -> Generator[Union[lzma.LZMAFile, io.TextIOWrapper[lzma.LZMAFile]], None, None]: """Ensure a LZMA-compressed file is downloaded and open a file inside it. :param subkeys: