Add type annotations to dulwich/cloud/gcs.py

dulwich/cloud/gcs.py

@@ -24,23 +24,26 @@
 
 import posixpath
 import tempfile
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, BinaryIO
 
 from ..object_store import BucketBasedObjectStore
-from ..pack import PACK_SPOOL_FILE_MAX_SIZE, Pack, PackData, load_pack_index_file
+from ..pack import (
+    PACK_SPOOL_FILE_MAX_SIZE,
+    Pack,
+    PackData,
+    PackIndex,
+    load_pack_index_file,
+)
+
+if TYPE_CHECKING:
+    from google.cloud.storage import Bucket
 
 # TODO(jelmer): For performance, read ranges?
 
 
 class GcsObjectStore(BucketBasedObjectStore):
-    """Object store implementation for Google Cloud Storage."""
-
-    def __init__(self, bucket, subpath="") -> None:
-        """Initialize GcsObjectStore.
-
-        Args:
-          bucket: GCS bucket instance
-          subpath: Subpath within the bucket
-        """
+    def __init__(self, bucket: "Bucket", subpath: str = "") -> None:
         super().__init__()
         self.bucket = bucket
         self.subpath = subpath
@@ -49,13 +52,13 @@ class GcsObjectStore(BucketBasedObjectStore):
         """Return string representation of GcsObjectStore."""
         return f"{type(self).__name__}({self.bucket!r}, subpath={self.subpath!r})"
 
-    def _remove_pack(self, name) -> None:
+    def _remove_pack(self, name: str) -> None:
         self.bucket.delete_blobs(
             [posixpath.join(self.subpath, name) + "." + ext for ext in ["pack", "idx"]]
         )
 
-    def _iter_pack_names(self):
-        packs = {}
+    def _iter_pack_names(self) -> Iterator[str]:
+        packs: dict[str, set[str]] = {}
         for blob in self.bucket.list_blobs(prefix=self.subpath):
             name, ext = posixpath.splitext(posixpath.basename(blob.name))
             packs.setdefault(name, set()).add(ext)
@@ -63,26 +66,26 @@ class GcsObjectStore(BucketBasedObjectStore):
             if exts == {".pack", ".idx"}:
                 yield name
 
-    def _load_pack_data(self, name):
+    def _load_pack_data(self, name: str) -> PackData:
         b = self.bucket.blob(posixpath.join(self.subpath, name + ".pack"))
         f = tempfile.SpooledTemporaryFile(max_size=PACK_SPOOL_FILE_MAX_SIZE)
         b.download_to_file(f)
         f.seek(0)
         return PackData(name + ".pack", f)
 
-    def _load_pack_index(self, name):
+    def _load_pack_index(self, name: str) -> PackIndex:
         b = self.bucket.blob(posixpath.join(self.subpath, name + ".idx"))
         f = tempfile.SpooledTemporaryFile(max_size=PACK_SPOOL_FILE_MAX_SIZE)
         b.download_to_file(f)
         f.seek(0)
         return load_pack_index_file(name + ".idx", f)
 
-    def _get_pack(self, name):
-        return Pack.from_lazy_objects(
+    def _get_pack(self, name: str) -> Pack:
+        return Pack.from_lazy_objects(  # type: ignore[no-untyped-call]
             lambda: self._load_pack_data(name), lambda: self._load_pack_index(name)
         )
 
-    def _upload_pack(self, basename, pack_file, index_file) -> None:
+    def _upload_pack(self, basename: str, pack_file: BinaryIO, index_file: BinaryIO) -> None:
         idxblob = self.bucket.blob(posixpath.join(self.subpath, basename + ".idx"))
         datablob = self.bucket.blob(posixpath.join(self.subpath, basename + ".pack"))
         idxblob.upload_from_file(index_file)