Skip to content

Commit

Permalink
Merge remote-tracking branch 'upstream/main' into feat-fsspecstore-bu…
Browse files Browse the repository at this point in the history
…lk-delete
  • Loading branch information
carshadi committed Jan 29, 2025
2 parents a0214e4 + 0c895d1 commit 75ee3a5
Show file tree
Hide file tree
Showing 35 changed files with 438 additions and 64 deletions.
4 changes: 2 additions & 2 deletions .github/labeler.yml
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
- needs release notes:
needs release notes:
- all:
- changed-files:
- any-glob-to-any-file: 'changes/*.rst'
- any-glob-to-any-file: '!changes/*.rst'
6 changes: 6 additions & 0 deletions .github/workflows/gpu_test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -64,3 +64,9 @@ jobs:
- name: Run Tests
run: |
hatch env run --env gputest.py${{ matrix.python-version }}-${{ matrix.numpy-version }}-${{ matrix.dependency-set }} run-coverage
- name: Upload coverage
uses: codecov/codecov-action@13ce06bfc6bbe3ecf90edbbf1bc32fe5978ca1d3 # v5.3.1
with:
token: ${{ secrets.CODECOV_TOKEN }}
verbose: true # optional (default = false)
2 changes: 1 addition & 1 deletion .github/workflows/releases.yml
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ jobs:
with:
name: releases
path: dist
- uses: pypa/[email protected].3
- uses: pypa/[email protected].4
with:
user: __token__
password: ${{ secrets.pypi_password }}
Expand Down
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ src/zarr/_version.py
data/*
src/fixture/
fixture/
junit.xml

.DS_Store
tests/.hypothesis
Expand Down
File renamed without changes.
1 change: 1 addition & 0 deletions changes/2693.bugfix.1.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Match the errors raised by read only stores in StoreTests.
1 change: 1 addition & 0 deletions changes/2693.bugfix.2.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Use stdout rather than stderr as the default stream for LoggingStore.
1 change: 1 addition & 0 deletions changes/2693.bugfix.3.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Ensure that ZipStore is open before getting or setting any values.
1 change: 1 addition & 0 deletions changes/2693.bugfix.4.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Update equality for LoggingStore and WrapperStore such that 'other' must also be a LoggingStore or WrapperStore respectively, rather than only checking the types of the stores they wrap.
1 change: 1 addition & 0 deletions changes/2693.feature.1.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Implemented open() for LoggingStore.
1 change: 1 addition & 0 deletions changes/2693.feature.2.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
LoggingStore is now a generic class.
3 changes: 3 additions & 0 deletions changes/2693.feature.3.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Change StoreTest's ``test_store_repr``, ``test_store_supports_writes``,
``test_store_supports_partial_writes``, and ``test_store_supports_listing``
to to be implemented using ``@abstractmethod``, rather raising ``NotImplementedError``.
1 change: 1 addition & 0 deletions changes/2693.feature.4.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Separate instantiating and opening a store in StoreTests.
1 change: 1 addition & 0 deletions changes/2693.feature.5.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add a test for using Stores as a context managers in StoreTests.
1 change: 1 addition & 0 deletions changes/2693.feature.6.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Test that a ValueError is raised for invalid byte range syntax in StoreTests.
1 change: 1 addition & 0 deletions changes/2693.feature.7.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Test getsize() and getsize_prefix() in StoreTests.
1 change: 1 addition & 0 deletions changes/2693.feature.8.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Test the error raised for invalid buffer arguments in StoreTests.
1 change: 1 addition & 0 deletions changes/2693.feature.9.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Test that data can be written to a store that's not yet open using the store.set method in StoreTests
2 changes: 2 additions & 0 deletions changes/2762.bugfix.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fixed ZipStore to make sure the correct attributes are saved when instances are pickled.
This fixes a previous bug that prevent using ZipStore with a ProcessPoolExecutor.
1 change: 1 addition & 0 deletions changes/2768.bugfix.1.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Updated the optional test dependencies to include ``botocore`` and ``fsspec``.
2 changes: 2 additions & 0 deletions changes/2768.bugfix.2.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fixed the fsspec tests to skip if ``botocore`` is not installed.
Previously they would have failed with an import error.
4 changes: 3 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -73,10 +73,12 @@ test = [
"coverage",
"pytest",
"pytest-cov",
'zarr[remote]',
"botocore",
"s3fs",
"moto[s3,server]",
"pytest-asyncio",
"pytest-accept",
"moto[s3,server]",
"requests",
"rich",
"mypy",
Expand Down
6 changes: 4 additions & 2 deletions src/zarr/abc/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,10 +176,10 @@ async def get(
Parameters
----------
key : str
prototype : BufferPrototype
The prototype of the output buffer. Stores may support a default buffer prototype.
byte_range : ByteRequest, optional
ByteRequest may be one of the following. If not provided, all data associated with the key is retrieved.
- RangeByteRequest(int, int): Request a specific range of bytes in the form (start, end). The end is exclusive. If the given range is zero-length or starts after the end of the object, an error will be returned. Additionally, if the range ends after the end of the object, the entire remainder of the object will be returned. Otherwise, the exact requested range will be returned.
- OffsetByteRequest(int): Request all bytes starting from a given byte offset. This is equivalent to bytes={int}- as an HTTP header.
- SuffixByteRequest(int): Request the last int bytes. Note that here, int is the size of the request, not the byte offset. This is equivalent to bytes=-{int} as an HTTP header.
Expand All @@ -200,6 +200,8 @@ async def get_partial_values(
Parameters
----------
prototype : BufferPrototype
The prototype of the output buffer. Stores may support a default buffer prototype.
key_ranges : Iterable[tuple[str, tuple[int | None, int | None]]]
Ordered set of key, range pairs, a key may occur multiple times with different ranges
Expand Down
7 changes: 6 additions & 1 deletion src/zarr/storage/_fsspec.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,15 @@
Store,
SuffixByteRequest,
)
from zarr.core.buffer import Buffer
from zarr.storage._common import _dereference_path

if TYPE_CHECKING:
from collections.abc import AsyncIterator, Iterable

from fsspec.asyn import AsyncFileSystem

from zarr.core.buffer import Buffer, BufferPrototype
from zarr.core.buffer import BufferPrototype
from zarr.core.common import BytesLike


Expand Down Expand Up @@ -266,6 +267,10 @@ async def set(
if not self._is_open:
await self._open()
self._check_writable()
if not isinstance(value, Buffer):
raise TypeError(
f"FsspecStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
)
path = _dereference_path(self.path, key)
# write data
if byte_range:
Expand Down
6 changes: 4 additions & 2 deletions src/zarr/storage/_local.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ def __init__(self, root: Path | str, *, read_only: bool = False) -> None:
root = Path(root)
if not isinstance(root, Path):
raise TypeError(
f'"root" must be a string or Path instance. Got an object with type {type(root)} instead.'
f"'root' must be a string or Path instance. Got an instance of {type(root)} instead."
)
self.root = root

Expand Down Expand Up @@ -169,7 +169,9 @@ async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
self._check_writable()
assert isinstance(key, str)
if not isinstance(value, Buffer):
raise TypeError("LocalStore.set(): `value` must a Buffer instance")
raise TypeError(
f"LocalStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
)
path = self.root / key
await asyncio.to_thread(_put, path, value, start=None, exclusive=exclusive)

Expand Down
26 changes: 18 additions & 8 deletions src/zarr/storage/_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,11 @@

import inspect
import logging
import sys
import time
from collections import defaultdict
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, Self, TypeVar

from zarr.abc.store import Store
from zarr.storage._wrapper import WrapperStore
Expand All @@ -18,8 +19,10 @@

counter: defaultdict[str, int]

T_Store = TypeVar("T_Store", bound=Store)

class LoggingStore(WrapperStore[Store]):

class LoggingStore(WrapperStore[T_Store]):
"""
Store wrapper that logs all calls to the wrapped store.
Expand All @@ -42,7 +45,7 @@ class LoggingStore(WrapperStore[Store]):

def __init__(
self,
store: Store,
store: T_Store,
log_level: str = "DEBUG",
log_handler: logging.Handler | None = None,
) -> None:
Expand All @@ -67,7 +70,7 @@ def _configure_logger(

def _default_handler(self) -> logging.Handler:
"""Define a default log handler"""
handler = logging.StreamHandler()
handler = logging.StreamHandler(stream=sys.stdout)
handler.setLevel(self.log_level)
handler.setFormatter(
logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
Expand All @@ -94,6 +97,14 @@ def log(self, hint: Any = "") -> Generator[None, None, None]:
end_time = time.time()
self.logger.info("Finished %s [%.2f s]", op, end_time - start_time)

@classmethod
async def open(cls: type[Self], store_cls: type[T_Store], *args: Any, **kwargs: Any) -> Self:
log_level = kwargs.pop("log_level", "DEBUG")
log_handler = kwargs.pop("log_handler", None)
store = store_cls(*args, **kwargs)
await store._open()
return cls(store=store, log_level=log_level, log_handler=log_handler)

@property
def supports_writes(self) -> bool:
with self.log():
Expand Down Expand Up @@ -126,8 +137,7 @@ def _is_open(self) -> bool:

@_is_open.setter
def _is_open(self, value: bool) -> None:
with self.log(value):
self._store._is_open = value
raise NotImplementedError("LoggingStore must be opened via the `_open` method")

async def _open(self) -> None:
with self.log():
Expand All @@ -151,11 +161,11 @@ def __str__(self) -> str:
return f"logging-{self._store}"

def __repr__(self) -> str:
return f"LoggingStore({repr(self._store)!r})"
return f"LoggingStore({self._store.__class__.__name__}, '{self._store}')"

def __eq__(self, other: object) -> bool:
with self.log(other):
return self._store == other
return type(self) is type(other) and self._store.__eq__(other._store) # type: ignore[attr-defined]

async def get(
self,
Expand Down
9 changes: 6 additions & 3 deletions src/zarr/storage/_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,9 @@ async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None
await self._ensure_open()
assert isinstance(key, str)
if not isinstance(value, Buffer):
raise TypeError(f"Expected Buffer. Got {type(value)}.")
raise TypeError(
f"MemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
)

if byte_range is not None:
buf = self._store_dict[key]
Expand Down Expand Up @@ -231,8 +233,9 @@ async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None
self._check_writable()
assert isinstance(key, str)
if not isinstance(value, Buffer):
raise TypeError(f"Expected Buffer. Got {type(value)}.")

raise TypeError(
f"GpuMemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
)
# Convert to gpu.Buffer
gpu_value = value if isinstance(value, gpu.Buffer) else gpu.Buffer.from_buffer(value)
await super().set(key, gpu_value, byte_range=byte_range)
16 changes: 15 additions & 1 deletion src/zarr/storage/_wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,14 @@ async def _ensure_open(self) -> None:
async def is_empty(self, prefix: str) -> bool:
return await self._store.is_empty(prefix)

@property
def _is_open(self) -> bool:
return self._store._is_open

@_is_open.setter
def _is_open(self, value: bool) -> None:
raise NotImplementedError("WrapperStore must be opened via the `_open` method")

async def clear(self) -> None:
return await self._store.clear()

Expand All @@ -67,7 +75,13 @@ def _check_writable(self) -> None:
return self._store._check_writable()

def __eq__(self, value: object) -> bool:
return type(self) is type(value) and self._store.__eq__(value)
return type(self) is type(value) and self._store.__eq__(value._store) # type: ignore[attr-defined]

def __str__(self) -> str:
return f"wrapping-{self._store}"

def __repr__(self) -> str:
return f"WrapperStore({self._store.__class__.__name__}, '{self._store}')"

async def get(
self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None
Expand Down
23 changes: 17 additions & 6 deletions src/zarr/storage/_zip.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,11 +107,14 @@ def _sync_open(self) -> None:
async def _open(self) -> None:
self._sync_open()

def __getstate__(self) -> tuple[Path, ZipStoreAccessModeLiteral, int, bool]:
return self.path, self._zmode, self.compression, self.allowZip64

def __setstate__(self, state: Any) -> None:
self.path, self._zmode, self.compression, self.allowZip64 = state
def __getstate__(self) -> dict[str, Any]:
state = self.__dict__
for attr in ["_zf", "_lock"]:
state.pop(attr, None)
return state

def __setstate__(self, state: dict[str, Any]) -> None:
self.__dict__ = state
self._is_open = False
self._sync_open()

Expand Down Expand Up @@ -146,6 +149,8 @@ def _get(
prototype: BufferPrototype,
byte_range: ByteRequest | None = None,
) -> Buffer | None:
if not self._is_open:
self._sync_open()
# docstring inherited
try:
with self._zf.open(key) as f: # will raise KeyError
Expand Down Expand Up @@ -190,6 +195,8 @@ async def get_partial_values(
return out

def _set(self, key: str, value: Buffer) -> None:
if not self._is_open:
self._sync_open()
# generally, this should be called inside a lock
keyinfo = zipfile.ZipInfo(filename=key, date_time=time.localtime(time.time())[:6])
keyinfo.compress_type = self.compression
Expand All @@ -203,9 +210,13 @@ def _set(self, key: str, value: Buffer) -> None:
async def set(self, key: str, value: Buffer) -> None:
# docstring inherited
self._check_writable()
if not self._is_open:
self._sync_open()
assert isinstance(key, str)
if not isinstance(value, Buffer):
raise TypeError("ZipStore.set(): `value` must a Buffer instance")
raise TypeError(
f"ZipStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
)
with self._lock:
self._set(key, value)

Expand Down
Loading

0 comments on commit 75ee3a5

Please sign in to comment.