Make ObjectStore a generic class (#3486)

Co-authored-by: Davis Bennett <[email protected]>

Author: dstansby

src/zarr/storage/_obstore.py

@@ -4,7 +4,7 @@
 import contextlib
 import pickle
 from collections import defaultdict
-from typing import TYPE_CHECKING, TypedDict
+from typing import TYPE_CHECKING, Generic, Self, TypedDict, TypeVar
 
 from zarr.abc.store import (
     ByteRequest,
@@ -34,7 +34,10 @@
 )
 
 
-class ObjectStore(Store):
+T_Store = TypeVar("T_Store", bound="_UpstreamObjectStore")
+
+
+class ObjectStore(Store, Generic[T_Store]):
     """
     Store that uses obstore for fast read/write from AWS, GCP, Azure.
 
@@ -51,7 +54,7 @@ class ObjectStore(Store):
     raise an issue with any comments/concerns about the store.
     """
 
-    store: _UpstreamObjectStore
+    store: T_Store
     """The underlying obstore instance."""
 
     def __eq__(self, value: object) -> bool:
@@ -61,15 +64,15 @@ def __eq__(self, value: object) -> bool:
         if not self.read_only == value.read_only:
             return False
 
-        return self.store == value.store
+        return self.store == value.store  # type: ignore[no-any-return]
 
-    def __init__(self, store: _UpstreamObjectStore, *, read_only: bool = False) -> None:
+    def __init__(self, store: T_Store, *, read_only: bool = False) -> None:
         if not store.__class__.__module__.startswith("obstore"):
             raise TypeError(f"expected ObjectStore class, got {store!r}")
         super().__init__(read_only=read_only)
         self.store = store
 
-    def with_read_only(self, read_only: bool = False) -> ObjectStore:
+    def with_read_only(self, read_only: bool = False) -> Self:
         # docstring inherited
         return type(self)(
             store=self.store,