Add overloads for lzma

cthoyt · Nov 15, 2024 · 5fef95b · 5fef95b
1 parent 24acde0
commit 5fef95b
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 6 deletions.
diff --git 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
@@ -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: