Fix typing errors in test_object

Author: dstansby

pyproject.toml

@@ -366,6 +366,7 @@ module = [
     "tests.test_store.test_memory",
     "tests.test_store.test_core",
     "tests.test_store.test_logging",
+    "tests.test_store.test_object",
 ]
 strict = false
 
@@ -375,7 +376,6 @@ strict = false
 module = [
     "tests.test_codecs.test_codecs",
     "tests.test_metadata.*",
-    "tests.test_store.test_object",
     "tests.test_store.test_stateful",
     "tests.test_store.test_wrapper",
     "tests.test_group",

src/zarr/storage/_obstore.py

@@ -5,7 +5,9 @@
 import pickle
 from collections import defaultdict
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any, TypedDict
+from typing import TYPE_CHECKING, Any, Generic, TypedDict, TypeVar
+
+from obstore.store import ObjectStore as _UpstreamObjectStore
 
 from zarr.abc.store import (
     ByteRequest,
@@ -22,11 +24,11 @@
     from typing import Any
 
     from obstore import ListResult, ListStream, ObjectMeta, OffsetRange, SuffixRange
-    from obstore.store import ObjectStore as _UpstreamObjectStore
 
     from zarr.core.buffer import Buffer, BufferPrototype
     from zarr.core.common import BytesLike
 
+
 __all__ = ["ObjectStore"]
 
 _ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = (
@@ -35,8 +37,10 @@
     NotADirectoryError,
 )
 
+T_Store = TypeVar("T_Store", bound=_UpstreamObjectStore)
 
-class ObjectStore(Store):
+
+class ObjectStore(Store, Generic[T_Store]):
     """
     Store that uses obstore for fast read/write from AWS, GCP, Azure.
 
@@ -53,19 +57,17 @@ 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:
-        if not isinstance(value, ObjectStore):
-            return False
-
-        if not self.read_only == value.read_only:
-            return False
-
-        return self.store == value.store
+        return bool(
+            isinstance(value, ObjectStore)
+            and self.read_only == value.read_only
+            and self.store == value.store
+        )
 
-    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)

tests/test_store/test_object.py

@@ -4,6 +4,8 @@
 import pytest
 
 obstore = pytest.importorskip("obstore")
+from pathlib import Path
+
 import pytest
 from hypothesis.stateful import (
     run_state_machine_as_test,
@@ -16,47 +18,43 @@
 from zarr.testing.store import StoreTests
 
 
-class TestObjectStore(StoreTests[ObjectStore, cpu.Buffer]):
+class TestObjectStore(StoreTests[ObjectStore[LocalStore], cpu.Buffer]):
     store_cls = ObjectStore
     buffer_cls = cpu.Buffer
 
     @pytest.fixture
-    def store_kwargs(self, tmpdir) -> dict[str, Any]:
-        store = LocalStore(prefix=tmpdir)
+    def store_kwargs(self, tmp_path: Path) -> dict[str, Any]:
+        store = LocalStore(prefix=str(tmp_path))
         return {"store": store, "read_only": False}
 
     @pytest.fixture
-    def store(self, store_kwargs: dict[str, str | bool]) -> ObjectStore:
+    async def store(self, store_kwargs: dict[str, Any]) -> ObjectStore[LocalStore]:
         return self.store_cls(**store_kwargs)
 
-    async def get(self, store: ObjectStore, key: str) -> Buffer:
-        assert isinstance(store.store, LocalStore)
+    async def get(self, store: ObjectStore[LocalStore], key: str) -> Buffer:
         new_local_store = LocalStore(prefix=store.store.prefix)
         return self.buffer_cls.from_bytes(obstore.get(new_local_store, key).bytes())
 
-    async def set(self, store: ObjectStore, key: str, value: Buffer) -> None:
-        assert isinstance(store.store, LocalStore)
+    async def set(self, store: ObjectStore[LocalStore], key: str, value: Buffer) -> None:
         new_local_store = LocalStore(prefix=store.store.prefix)
         obstore.put(new_local_store, key, value.to_bytes())
 
-    def test_store_repr(self, store: ObjectStore) -> None:
+    def test_store_repr(self, store: ObjectStore[LocalStore]) -> None:
         from fnmatch import fnmatch
 
         pattern = "ObjectStore(object_store://LocalStore(*))"
         assert fnmatch(f"{store!r}", pattern)
 
-    def test_store_supports_writes(self, store: ObjectStore) -> None:
+    def test_store_supports_writes(self, store: ObjectStore[LocalStore]) -> None:
         assert store.supports_writes
 
-    async def test_store_supports_partial_writes(self, store: ObjectStore) -> None:
+    def test_store_supports_partial_writes(self, store: ObjectStore[LocalStore]) -> None:
         assert not store.supports_partial_writes
-        with pytest.raises(NotImplementedError):
-            await store.set_partial_values([("foo", 0, b"\x01\x02\x03\x04")])
 
-    def test_store_supports_listing(self, store: ObjectStore) -> None:
+    def test_store_supports_listing(self, store: ObjectStore[LocalStore]) -> None:
         assert store.supports_listing
 
-    def test_store_equal(self, store: ObjectStore) -> None:
+    def test_store_equal(self, store: ObjectStore[LocalStore]) -> None:
         """Test store equality"""
         # Test equality against a different instance type
         assert store != 0
@@ -73,14 +71,14 @@ def test_store_equal(self, store: ObjectStore) -> None:
     def test_store_init_raises(self) -> None:
         """Test __init__ raises appropriate error for improper store type"""
         with pytest.raises(TypeError):
-            ObjectStore("path/to/store")
+            ObjectStore("path/to/store")  # type: ignore[type-var]
 
 
 @pytest.mark.slow_hypothesis
-def test_zarr_hierarchy():
+def test_zarr_hierarchy() -> None:
     sync_store = ObjectStore(MemoryStore())
 
     def mk_test_instance_sync() -> ZarrHierarchyStateMachine:
         return ZarrHierarchyStateMachine(sync_store)
 
-    run_state_machine_as_test(mk_test_instance_sync)
+    run_state_machine_as_test(mk_test_instance_sync)  # type ignore[no-untyped-call]