diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3ecbda5..624f06a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -32,3 +32,5 @@ repos: # - types-lxml - ome-types - pydantic + - fsspec + diff --git a/docs/index.md b/docs/index.md index 379a193..10b3d6f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -44,6 +44,16 @@ install with support for these files use the `legacy` extra: pip install nd2[legacy] ``` +### With remote nd2 file support + +If you would like to be able to read nd2 files from remote sources (e.g. S3, HTTP), +then fsspec and/or s3fs are required. +To install with support for remote files use the `remote` extra: + +```sh +pip install nd2[remote] +``` + ### Faster XML parsing Much of the metadata in the file stored as XML. If found in the environment, diff --git a/mkdocs.yml b/mkdocs.yml index 07c81e7..c8eadc6 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -64,6 +64,7 @@ plugins: - https://docs.xarray.dev/en/stable/objects.inv - https://docs.dask.org/en/stable/objects.inv - https://ome-types.readthedocs.io/en/latest/objects.inv + - https://filesystem-spec.readthedocs.io/en/latest/objects.inv options: docstring_section_style: list # or "table" diff --git a/pyproject.toml b/pyproject.toml index 6a15bd2..c53a746 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,16 +47,18 @@ ome-zarr-tensorstore = [ "pydantic>=2.11.0", "yaozarrs[write-tensorstore]>=0.2.0; python_version >= '3.10'", ] +remote = ["fsspec>=2024.2.0", "s3fs>=2024.2.0"] [dependency-groups] test = [ - "nd2[tiff]", + "nd2[tiff,remote]", "dask[array]>=2023.5.0", "imagecodecs>=2023.1.23", # older versions no longer ship wheels for these pythons "imagecodecs>=2024.6.1; python_version >= '3.12'", "imagecodecs>=2024.12.30; python_version >= '3.13'", "lxml>=5.3.0", + "moto[s3,server]>=5.1.0", "psutil>=5.9.8", "pytest-codspeed>=4.1.1", "pytest-cov>=4.0.0", @@ -68,6 +70,7 @@ test = [ dev = [ { include-group = "test" }, { include-group = "docs" }, + "boto3-stubs[s3]>=1.42.52", "ipython>=8.0.0", "mypy>=1.18.2", "pdbpp>=0.11.7; sys_platform != 'win32'", @@ -75,6 +78,7 @@ dev = [ "pytest-xdist>=3.8.0", "rich>=14.2.0", "ruff>=0.14.0", + "s3fs>=2025.10.0", "types-lxml>=2025.8.25", ] docs = [ @@ -158,6 +162,7 @@ filterwarnings = [ "ignore:::xarray", "ignore:Accessing the 'model_fields' attribute::", "ignore:Failing to pass a value:DeprecationWarning", # xsdata + "ignore:Boto3 will no longer support Python 3.9" ] # https://mypy.readthedocs.io/en/stable/config_file.html diff --git a/src/nd2/_nd2file.py b/src/nd2/_nd2file.py index 1835a54..9342016 100644 --- a/src/nd2/_nd2file.py +++ b/src/nd2/_nd2file.py @@ -5,6 +5,7 @@ import threading import warnings from itertools import product +from pathlib import Path from types import MappingProxyType from typing import TYPE_CHECKING, Callable, cast, overload @@ -13,7 +14,7 @@ from nd2 import _util from ._readers.protocol import ND2Reader -from ._util import AXIS, is_supported_file +from ._util import AXIS, is_fsspec_url, is_supported_file try: from functools import cached_property @@ -24,7 +25,6 @@ if TYPE_CHECKING: from collections.abc import Mapping, Sequence, Sized from os import PathLike - from pathlib import Path from typing import Any, Literal, SupportsInt import dask.array @@ -83,8 +83,14 @@ class ND2File: Parameters ---------- - path : Path | str - Filename of an nd2 file. + path : Path | str | ReadSeekBinary + Filename, Path, or URL of an nd2 file. May also be any file-like object that + supports binary read and seek, including those returned by `fsspec` filesystems: + + - `closed(self) -> bool: ...` + - `read(self, size: int | None = ..., /) -> bytes: ...` + - `seek(self, offset: int, whence: int = ..., /) -> int: ...` + - `close(self) -> None: ...` validate_frames : bool Whether to verify (and attempt to fix) frames whose positions have been shifted relative to the predicted offset (i.e. in a corrupted file). @@ -93,6 +99,8 @@ class ND2File: search_window : int When validate_frames is true, this is the search window (in KB) that will be used to try to find the actual chunk position. by default 100 KB + storage_options : dict, optional + Extra kwargs passed to [`fsspec.core.url_to_fs`][] when opening remote URLs. """ def __init__( @@ -101,12 +109,18 @@ def __init__( *, validate_frames: bool = False, search_window: int = 100, + storage_options: dict[str, Any] | None = None, ) -> None: self._error_radius: int | None = ( search_window * 1000 if validate_frames else None ) - self._rdr = ND2Reader.create(path, self._error_radius) - self._path = self._rdr._path + self._storage_options = storage_options + self._rdr = ND2Reader.create(path, self._error_radius, storage_options) + self._path: str | Path | None = self._rdr._path + # For URL inputs the reader may not extract the full URL from the handle; + # preserve the original URL string so pickling round-trips correctly. + if is_fsspec_url(path): + self._path = path self._lock = threading.RLock() @staticmethod @@ -142,7 +156,7 @@ def version(self) -> tuple[int, ...]: @property def path(self) -> str: """Path of the image.""" - return str(self._path) + return "" if self._path is None else str(self._path) @property def is_legacy(self) -> bool: @@ -196,7 +210,9 @@ def __enter__(self) -> ND2File: def __del__(self) -> None: """Delete file handle on garbage collection.""" # if it came in as an open file handle, it's ok to remain open after deletion - if not getattr(self, "closed", True) and not self._rdr._was_open: + if not hasattr(self, "_rdr"): + return + if not self.closed and not self._rdr._was_open: # this stack inspection is a hack to avoid an unnecessary warning/closure. # when using the to_dask() method, calling dask map_blocks will greedily # pickle/unpickle the object. @@ -218,6 +234,10 @@ def __exit__(self, *_: Any) -> None: def __getstate__(self) -> dict[str, Any]: """Return state for pickling.""" + if self._path is None: + raise TypeError( + "Cannot pickle ND2File created from an unnamed file-like object" + ) state = self.__dict__.copy() del state["_rdr"] del state["_lock"] @@ -230,7 +250,11 @@ def __setstate__(self, d: dict[str, Any]) -> None: _was_closed = d.pop("_closed", False) self.__dict__ = d self._lock = threading.RLock() - self._rdr = ND2Reader.create(self._path, self._error_radius) + self._rdr = ND2Reader.create( + self._path, # type: ignore + self._error_radius, + cast("dict | None", self.__dict__.get("_storage_options")), + ) if _was_closed: self._rdr.close() @@ -1352,7 +1376,8 @@ def __repr__(self) -> str: """Return a string representation of the ND2File.""" try: details = " (closed)" if self.closed else f" {self.dtype}: {self.sizes!r}" - extra = f": {self._path.name!r}{details}" + path_name = self._path.name if isinstance(self._path, Path) else self._path + extra = f": {path_name!r}{details}" except Exception: extra = "" return f"" @@ -1453,6 +1478,7 @@ def imread( dask: Literal[False] = ..., xarray: Literal[False] = ..., validate_frames: bool = ..., + storage_options: dict | None = ..., ) -> np.ndarray: ... @@ -1463,6 +1489,7 @@ def imread( dask: bool = ..., xarray: Literal[True], validate_frames: bool = ..., + storage_options: dict | None = ..., ) -> xr.DataArray: ... @@ -1473,6 +1500,7 @@ def imread( dask: Literal[True], xarray: Literal[False] = ..., validate_frames: bool = ..., + storage_options: dict | None = ..., ) -> dask.array.core.Array: ... @@ -1482,13 +1510,15 @@ def imread( dask: bool = False, xarray: bool = False, validate_frames: bool = False, + storage_options: dict | None = None, ) -> np.ndarray | xr.DataArray | dask.array.core.Array: """Open `file`, return requested array type, and close `file`. Parameters ---------- file : Path | str - Filepath (`str`) or `Path` object to ND2 file. + Filepath (`str`) or `Path` object to ND2 file. May also be a remote + URL string (e.g. ``"s3://bucket/key.nd2"``) when fsspec is installed. dask : bool If `True`, returns a (delayed) `dask.array.Array`. This will avoid reading any data from disk until specifically requested by using `.compute()` or @@ -1504,13 +1534,18 @@ def imread( shifted relative to the predicted offset (i.e. in a corrupted file). This comes at a slight performance penalty at file open, but may "rescue" some corrupt files. by default False. + storage_options : dict, optional + Extra kwargs forwarded to fsspec when opening a remote URL, e.g. + ``{"anon": True}`` for public S3 buckets. Returns ------- Union[np.ndarray, dask.array.Array, xarray.DataArray] Array subclass, depending on arguments used. """ - with ND2File(file, validate_frames=validate_frames) as nd2: + with ND2File( + file, validate_frames=validate_frames, storage_options=storage_options + ) as nd2: if xarray: return nd2.to_xarray(delayed=dask) elif dask: diff --git a/src/nd2/_ome.py b/src/nd2/_ome.py index a75c5bc..3c83524 100644 --- a/src/nd2/_ome.py +++ b/src/nd2/_ome.py @@ -246,6 +246,7 @@ def nd2_ome_metadata( if include_unstructured: all_meta = m.MapAnnotation( + id="Annotation:0", description="ND2 unstructured metadata, encoded as a JSON string. " "Each key in this MapAnnotation is the name of a metadata chunk found in " "the ND2 file, and the value is the JSON-encoded data for that chunk.", diff --git a/src/nd2/_parse/_chunk_decode.py b/src/nd2/_parse/_chunk_decode.py index b159a82..c8d836d 100644 --- a/src/nd2/_parse/_chunk_decode.py +++ b/src/nd2/_parse/_chunk_decode.py @@ -4,20 +4,21 @@ import mmap import struct -from contextlib import contextmanager, nullcontext -from pathlib import Path -from typing import TYPE_CHECKING, BinaryIO, cast +from contextlib import AbstractContextManager, contextmanager, nullcontext +from os import PathLike as OSPathLike +from typing import TYPE_CHECKING, cast import numpy as np if TYPE_CHECKING: from collections.abc import Iterator - from contextlib import AbstractContextManager from os import PathLike from typing import Final from numpy.typing import DTypeLike + from nd2._util import ReadSeekBinary + StrOrBytesPath = str | bytes | PathLike[str] | PathLike[bytes] StartFileChunk = tuple[int, int, int, bytes, bytes] @@ -68,12 +69,12 @@ # uint64_t offset -def get_version(fh: BinaryIO | StrOrBytesPath) -> tuple[int, int]: +def get_version(fh: ReadSeekBinary | StrOrBytesPath) -> tuple[int, int]: """Get the version of the ND2 file or raise an exception. Parameters ---------- - fh : BinaryIO | str | bytes | Path + fh : ReadSeekBinary | str | bytes | Path The file handle or path to the ND2 file. Returns @@ -87,14 +88,14 @@ def get_version(fh: BinaryIO | StrOrBytesPath) -> tuple[int, int]: If the file is not a valid ND2 file or the header chunk is corrupt. """ if hasattr(fh, "read"): - ctx: AbstractContextManager[BinaryIO] = nullcontext(cast("BinaryIO", fh)) + ctx: AbstractContextManager = nullcontext(cast("ReadSeekBinary", fh)) else: ctx = open(fh, "rb") - with ctx as fh: - fh.seek(0) - fname = str(fh.name) - chunk = START_FILE_CHUNK.unpack(fh.read(START_FILE_CHUNK.size)) + with ctx as handle: + handle.seek(0) + fname = str(getattr(handle, "name", "")) + chunk = START_FILE_CHUNK.unpack(handle.read(START_FILE_CHUNK.size)) magic, name_length, data_length, name, data = cast("StartFileChunk", chunk) @@ -112,7 +113,7 @@ def get_version(fh: BinaryIO | StrOrBytesPath) -> tuple[int, int]: return (int(chr(data[3])), int(chr(data[5]))) -def get_chunkmap(fh: BinaryIO, error_radius: int | None = None) -> ChunkMap: +def get_chunkmap(fh: ReadSeekBinary, error_radius: int | None = None) -> ChunkMap: """Read the map of the chunks at the end of an ND2 file. A Chunkmap is mapping of chunk names (bytes) to (offset, size) pairs. @@ -128,7 +129,7 @@ def get_chunkmap(fh: BinaryIO, error_radius: int | None = None) -> ChunkMap: Parameters ---------- - fh : BinaryIO + fh : ReadSeekBinary An open nd2 file. File is assumed to be a valid ND2 file. (use `get_version`) error_radius : int, optional If b"ND2 FILEMAP SIGNATURE NAME 0001!" is not found at expected location and @@ -149,7 +150,10 @@ def get_chunkmap(fh: BinaryIO, error_radius: int | None = None) -> ChunkMap: fh.seek(-40, 2) sig, location = SIG_CHUNKMAP_LOC.unpack(fh.read(SIG_CHUNKMAP_LOC.size)) if sig != ND2_CHUNKMAP_SIGNATURE: # pragma: no cover - raise ValueError(f"Invalid ChunkMap signature {sig!r} in file {fh.name!r}") + raise ValueError( + f"Invalid ChunkMap signature {sig!r} in file " + f"{getattr(fh, 'name', '')!r}" + ) # get all of the data in the chunkmap chunkmap_data = _robustly_read_named_chunk( @@ -182,7 +186,7 @@ def get_chunkmap(fh: BinaryIO, error_radius: int | None = None) -> ChunkMap: def read_nd2_chunk( - fh: BinaryIO, start_position: int, expect_name: bytes | None = None + fh: ReadSeekBinary, start_position: int, expect_name: bytes | None = None ) -> bytes: """Read a single chunk in an ND2 file at `start_position`. @@ -197,7 +201,7 @@ def read_nd2_chunk( Parameters ---------- - fh : BinaryIO + fh : ReadSeekBinary An open nd2 file. File is assumed to be a valid ND2 file. (use `get_version`) start_position : int The position in the file to start reading the chunk. @@ -235,7 +239,7 @@ def read_nd2_chunk( def _robustly_read_named_chunk( - fh: BinaryIO, + fh: ReadSeekBinary, start_position: int, expect_name: bytes = ND2_FILEMAP_SIGNATURE, search_radius: int | None = None, @@ -248,7 +252,7 @@ def _robustly_read_named_chunk( Parameters ---------- - fh : BinaryIO + fh : ReadSeekBinary An open nd2 file. File is assumed to be a valid ND2 file. start_position : int The position in the file to start reading the chunk. @@ -261,8 +265,9 @@ def _robustly_read_named_chunk( try: return read_nd2_chunk(fh, start_position, expect_name=expect_name) except ValueError as e: + file_name = getattr(fh, "name", "") err_msg = ( - f"File {fh.name!r} appears to be corrupt. Expected " + f"File {file_name!r} appears to be corrupt. Expected " f"{expect_name!r} at position " f"{start_position} but did not find it." ) @@ -282,7 +287,7 @@ def _robustly_read_named_chunk( raise ValueError(err_msg) from e -def iter_chunks(handle: BinaryIO) -> Iterator[tuple[str, int, int]]: +def iter_chunks(handle: ReadSeekBinary) -> Iterator[tuple[str, int, int]]: file_size = handle.seek(0, 2) handle.seek(0) pos = 0 @@ -304,7 +309,7 @@ def iter_chunks(handle: BinaryIO) -> Iterator[tuple[str, int, int]]: def rescue_nd2( - handle: BinaryIO | str, + handle: ReadSeekBinary | StrOrBytesPath, frame_shape: tuple[int, ...] = (), dtype: DTypeLike = "uint16", max_iters: int | None = None, @@ -321,7 +326,7 @@ def rescue_nd2( Parameters ---------- - handle : BinaryIO | str + handle : ReadSeekBinary | str | bytes | PathLike Filepath string, or binary file handle (For example `handle = open('some.nd2', 'rb')`) frame_shape : Tuple[int, ...], optional @@ -356,7 +361,10 @@ def rescue_nd2( """ dtype = np.dtype(dtype) with ensure_handle(handle) as _fh: - mm = mmap.mmap(_fh.fileno(), 0, access=mmap.ACCESS_READ) + fileno = getattr(_fh, "fileno", None) + if not callable(fileno): + raise TypeError("rescue_nd2 requires a file handle with fileno()") + mm = mmap.mmap(fileno(), 0, access=mmap.ACCESS_READ) offset = 0 iters = 0 @@ -407,11 +415,15 @@ def rescue_nd2( @contextmanager -def ensure_handle(obj: str | BinaryIO) -> Iterator[BinaryIO]: - fh = open(obj, "rb") if isinstance(obj, (str, bytes, Path)) else obj +def ensure_handle(obj: StrOrBytesPath | ReadSeekBinary) -> Iterator[ReadSeekBinary]: + if isinstance(obj, (str, bytes, OSPathLike)): + opened_here = True + fh = cast("ReadSeekBinary", open(obj, "rb")) + else: + opened_here = False + fh = obj try: yield fh finally: - # close it if we were the one to open it - if not hasattr(obj, "fileno"): + if opened_here: fh.close() diff --git a/src/nd2/_readers/_legacy/legacy_reader.py b/src/nd2/_readers/_legacy/legacy_reader.py index f47a19f..d60d1e8 100644 --- a/src/nd2/_readers/_legacy/legacy_reader.py +++ b/src/nd2/_readers/_legacy/legacy_reader.py @@ -23,9 +23,9 @@ if TYPE_CHECKING: from collections.abc import Mapping - from typing import Any, BinaryIO, TypedDict + from typing import Any, TypedDict - from nd2._util import FileOrBinaryIO + from nd2._util import FileOrBinaryIO, ReadSeekBinary class RawExperimentLoop(TypedDict, total=False): Type: int @@ -365,9 +365,10 @@ def read_frame(self, index: int) -> np.ndarray: try: from imagecodecs import jpeg2k_decode except ModuleNotFoundError as e: # pragma: no cover + file_name = getattr(self._fh, "name", "") raise ModuleNotFoundError( f"{e}\n" - f"Reading legacy format nd2 {self._fh.name!r} requires imagecodecs.\n" + f"Reading legacy format nd2 {file_name!r} requires imagecodecs.\n" "Please install with `pip install imagecodecs`." ) from e @@ -415,7 +416,7 @@ def header(self) -> dict: pos = self.chunkmap[b"jp2h"][0] except (KeyError, IndexError) as e: # pragma: no cover raise KeyError("No valid jp2h header found in file") from e - fh = cast("BinaryIO", self._fh) + fh = cast("ReadSeekBinary", self._fh) fh.seek(pos + I4s.size + 4) # 4 bytes for "label" if fh.read(4) != b"ihdr": raise KeyError("No valid ihdr header found in jp2h header") @@ -438,7 +439,7 @@ def events(self, orient: str, null_value: Any) -> list | Mapping: return [] if orient == "records" else {} -def legacy_nd2_chunkmap(fh: BinaryIO) -> dict[bytes, list[int]]: +def legacy_nd2_chunkmap(fh: ReadSeekBinary) -> dict[bytes, list[int]]: fh.seek(-40, 2) sig, map_start = struct.unpack("<32sQ", fh.read()) if sig != b"LABORATORY IMAGING ND BOX MAP 00": # pragma: no cover diff --git a/src/nd2/_readers/_modern/modern_reader.py b/src/nd2/_readers/_modern/modern_reader.py index 7655178..a0b0658 100644 --- a/src/nd2/_readers/_modern/modern_reader.py +++ b/src/nd2/_readers/_modern/modern_reader.py @@ -1,6 +1,5 @@ from __future__ import annotations -import os import warnings import zlib from contextlib import suppress @@ -180,9 +179,6 @@ def _cached_global_metadata(self) -> GlobalMetadata: exp_loops=self.experiment(), text_info=self.text_info(), ) - if self._global_metadata["time"]["absoluteJulianDayNumber"] < 1: - julian_day = os.stat(self._path).st_ctime / 86400.0 + 2440587.5 - self._global_metadata["time"]["absoluteJulianDayNumber"] = julian_day return self._global_metadata @@ -331,18 +327,21 @@ def read_frame(self, index: int) -> np.ndarray: if self.attributes().compressionType == "lossless": return self._read_compressed_frame(index) - try: - return np.ndarray( - shape=self._actual_frame_shape(), - dtype=self._dtype(), - buffer=self._mmap, - offset=offset, - strides=self._strides, - ) - except TypeError: - # If the chunkmap is wrong, and the mmap isn't long enough - # for the requested offset & size, a TypeError is raised. - return self._missing_frame(index) + if self._mmap is not None: + try: + return np.ndarray( + shape=self._actual_frame_shape(), + dtype=self._dtype(), + buffer=self._mmap, + offset=offset, + strides=self._strides, + ) + except TypeError: + # If the chunkmap is wrong, and the mmap isn't long enough + # for the requested offset & size, a TypeError is raised. + return self._missing_frame(index) + + return self._read_frame_bytes(offset) def _read_compressed_frame(self, index: int) -> np.ndarray: ch = self._load_chunk(f"ImageDataSeq|{index}!".encode()) @@ -353,6 +352,29 @@ def _read_compressed_frame(self, index: int) -> np.ndarray: strides=self._strides, ) + def _read_frame_bytes(self, offset: int) -> np.ndarray: + """Read a frame via seek/read (fallback when mmap is unavailable).""" + if self._fh is None: # pragma: no cover + raise ValueError("Attempt to read from closed nd2 file") + + shape = self._actual_frame_shape() + dtype = self._dtype() + if self._strides is not None: + nbytes = shape[0] * (self.attributes().widthBytes or 0) + else: + nbytes = int(np.prod(shape)) * dtype.itemsize + self._fh.seek(offset) + data = self._fh.read(nbytes) + if self._strides is not None: + arr = np.ndarray( + shape=shape, + dtype=dtype, + buffer=data, + strides=self._strides, + ) + return arr.copy() + return np.frombuffer(data, dtype=dtype).reshape(shape) + def _missing_frame(self, index: int = 0) -> np.ndarray: # TODO: add other modes for filling missing data return np.zeros(self._raw_frame_shape(), self._dtype()) diff --git a/src/nd2/_readers/protocol.py b/src/nd2/_readers/protocol.py index 5a379f2..8102577 100644 --- a/src/nd2/_readers/protocol.py +++ b/src/nd2/_readers/protocol.py @@ -3,21 +3,21 @@ import abc import mmap import warnings -from contextlib import nullcontext +from contextlib import AbstractContextManager, nullcontext, suppress from pathlib import Path -from typing import TYPE_CHECKING, BinaryIO, cast +from typing import TYPE_CHECKING, Any, cast from nd2._parse._chunk_decode import get_version +from nd2._util import is_fsspec_url, is_read_seek_binary, open_fsspec_url if TYPE_CHECKING: from collections.abc import Mapping, Sequence - from contextlib import AbstractContextManager - from typing import Any, Literal + from typing import Literal import numpy as np from nd2._binary import BinaryLayers - from nd2._util import FileOrBinaryIO + from nd2._util import FileOrBinaryIO, ReadSeekBinary from nd2.jobs.types import JobsDict from nd2.structures import ( ROI, @@ -41,6 +41,7 @@ def create( cls, path: FileOrBinaryIO, error_radius: int | None = None, + storage_options: dict | None = None, ) -> ND2Reader: """Create an ND2Reader for the given path, using the appropriate subclass. @@ -52,44 +53,70 @@ def create( If b"ND2 FILEMAP SIGNATURE NAME 0001!" is not found at expected location and `error_radius` is not None, then an area of +/- `error_radius` bytes will be searched for the signature. + storage_options : dict, optional + Extra kwargs passed to fsspec when opening remote URLs. """ from nd2._readers import LegacyReader, ModernReader - if hasattr(path, "read"): - path = cast("BinaryIO", path) - if "b" not in path.mode: + is_url = is_fsspec_url(path) + opened_here = False + ctx: AbstractContextManager + if is_file_handle := is_read_seek_binary(path): + mode = getattr(path, "mode", "b") + if isinstance(mode, str) and "b" not in mode: raise ValueError( "File handles passed to ND2File must be in binary mode" ) - ctx: AbstractContextManager[BinaryIO] = nullcontext(path) + ctx = nullcontext(path) + elif is_url: + fh = open_fsspec_url(str(path), storage_options=storage_options) + opened_here = True + ctx = nullcontext(fh) else: - path = Path(path).expanduser().absolute() + path = Path(cast("str | Path", path)).expanduser().absolute() ctx = open(path, "rb") with ctx as fh: - fname = fh.name + fname = getattr(fh, "name", "") fh.seek(0) magic_num = fh.read(4) for subcls in (ModernReader, LegacyReader): if magic_num == subcls.HEADER_MAGIC: - return subcls(path, error_radius=error_radius) + # For URL/file-like cases pass the open handle; for local paths + # pass the Path so the reader can reopen it as needed. + effective_path = fh if (is_url or is_file_handle) else path + try: + return subcls(effective_path, error_radius=error_radius) + except Exception: + if opened_here: + fh.close() + raise + if opened_here: + fh.close() raise OSError( - f"file {fname} not recognized as ND2. First 4 bytes: {magic_num!r}" + f"file {fname!r} not recognized as ND2. First 4 bytes: {magic_num!r}" ) - def __init__(self, path: FileOrBinaryIO, error_radius: int | None = None) -> None: + def __init__(self, obj: FileOrBinaryIO, error_radius: int | None = None) -> None: self._chunkmap: dict | None = None + self._version: tuple[int, int] | None = None self._mmap: mmap.mmap | None = None - if hasattr(path, "read"): - self._fh: BinaryIO | None = cast("BinaryIO", path) - self._was_open = not self._fh.closed - self._path: Path = Path(self._fh.name) - self._mmap = mmap.mmap(self._fh.fileno(), 0, access=mmap.ACCESS_READ) + self._fh: ReadSeekBinary | None + self._path: str | Path | None + if is_read_seek_binary(obj): + self._fh = obj + self._was_open = not obj.closed + name = getattr(obj, "full_name", None) or getattr(obj, "name", None) + self._path = name if isinstance(name, str) else None + with suppress(Exception): + # remote/non-fileno file-likes: mmap not available + if (fileno := getattr(self._fh, "fileno", None)) and callable(fileno): + self._mmap = mmap.mmap(fileno(), 0, access=mmap.ACCESS_READ) else: self._was_open = False - self._path = Path(path) + self._path = Path(cast("str | Path", obj)) self._fh = None self._error_radius: int | None = error_radius self.open() @@ -101,8 +128,13 @@ def is_legacy(self) -> bool: def open(self) -> None: """Open the file handle.""" if self._fh is None or self._fh.closed: + if not isinstance(self._path, Path): + raise RuntimeError( + "Cannot reopen a remote/file-like ND2 source after closing" + ) self._fh = open(self._path, "rb") - self._mmap = mmap.mmap(self._fh.fileno(), 0, access=mmap.ACCESS_READ) + fh = cast("Any", self._fh) + self._mmap = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) def close(self) -> None: """Close the file handle.""" @@ -128,7 +160,16 @@ def __exit__(self, *_: Any) -> None: def version(self) -> tuple[int, int]: """Return the file format version as a tuple of ints.""" - return get_version(self._fh or self._path) + if self._version is None: + if self._fh is not None: + self._version = get_version(self._fh) + elif self._path is not None: + self._version = get_version(self._path) + else: + raise RuntimeError( + "Cannot determine version without an open file handle" + ) + return self._version def rois(self) -> list[ROI]: """Return ROIs in the file.""" diff --git a/src/nd2/_util.py b/src/nd2/_util.py index 2a35421..4df784f 100644 --- a/src/nd2/_util.py +++ b/src/nd2/_util.py @@ -1,59 +1,112 @@ from __future__ import annotations +import io import math import re from contextlib import suppress from datetime import datetime, timezone from itertools import product -from typing import TYPE_CHECKING, BinaryIO, NamedTuple, cast +from typing import TYPE_CHECKING, NamedTuple, cast if TYPE_CHECKING: from collections.abc import Mapping, Sequence from os import PathLike - from typing import Any, Callable, ClassVar, Final, Union + from typing import Any, ClassVar, Final, Protocol, TypeAlias, Union + + from typing_extensions import TypeGuard from nd2.structures import ExpLoop - StrOrPath = Union[str, PathLike] - FileOrBinaryIO = Union[StrOrPath, BinaryIO] + class ReadSeekBinary(Protocol): + @property + def closed(self) -> bool: ... + def read(self, size: int | None = -1, /) -> bytes: ... + def seek(self, offset: int, whence: int = 0, /) -> int: ... + def close(self) -> None: ... + + StrOrPath: TypeAlias = Union[str, PathLike] + FileOrBinaryIO: TypeAlias = Union[StrOrPath, ReadSeekBinary] - ListOfDicts = list[dict[str, Any]] - DictOfLists = Mapping[str, Sequence[Any]] - DictOfDicts = Mapping[str, dict[int, Any]] + ListOfDicts: TypeAlias = list[dict[str, Any]] + DictOfLists: TypeAlias = Mapping[str, Sequence[Any]] + DictOfDicts: TypeAlias = Mapping[str, dict[int, Any]] NEW_HEADER_MAGIC = b"\xda\xce\xbe\n" OLD_HEADER_MAGIC = b"\x00\x00\x00\x0c" VERSION = re.compile(r"^ND2 FILE SIGNATURE CHUNK NAME01!Ver([\d\.]+)$") -def _open_binary(path: StrOrPath) -> BinaryIO: - return open(path, "rb") +def is_fsspec_url(path: Any) -> TypeGuard[str]: + """True if `path` is a string with a remote URL scheme (e.g. 's3://').""" + if not isinstance(path, str): + return False + # RFC 3986 scheme syntax + return bool(re.match(r"^[a-zA-Z][a-zA-Z0-9+\-.]+://", path)) + + +def open_fsspec_url( + url: str, + *, + small_file_limit: int = 32 * 1024 * 1024, # 32 MB + block_size: int = 512 * 1024, # 512 KB + storage_options: dict | None = None, +) -> ReadSeekBinary: + """Open a remote nd2 URL. + + Files up to `small_file_limit` are fetched eagerly into a BytesIO object. + ND2 access is heavily random (magic bytes at start, chunkmap at end, + scattered frame offsets), so keeping small files fully in memory avoids + repeated range requests. + + For larger files a 512 KB `block_size` is used rather than fsspec's 5 MB + default. Benchmarking shows this cuts open + chunkmap latency ~2.5x: + each chunkmap seek fetches 512 KB instead of 5 MB, and scattered metadata + chunks typically fit within a single block at this size. + """ + try: + import fsspec + except ImportError as e: + raise ImportError( + "fsspec is required to open remote URLs. Install with: pip install fsspec" + ) from e + + sopts = storage_options or {} + fs, fpath = fsspec.url_to_fs(url, **sopts) + size = fs.info(fpath).get("size") + + if size is not None and size <= small_file_limit: + return io.BytesIO(fs.cat_file(fpath)) + + return cast("ReadSeekBinary", fs.open(fpath, "rb", block_size=block_size)) + + +def is_read_seek_binary(obj: object) -> TypeGuard[ReadSeekBinary]: + return ( + hasattr(obj, "read") + and hasattr(obj, "seek") + and hasattr(obj, "close") + and hasattr(obj, "closed") + ) -def is_supported_file( - path: FileOrBinaryIO, - open_: Callable[[StrOrPath], BinaryIO] = _open_binary, -) -> bool: +def is_supported_file(path: FileOrBinaryIO) -> bool: """Return `True` if `path` can be opened as an nd2 file. Parameters ---------- path : Union[str, bytes, PathLike] A path to query - open_ : Callable[[StrOrBytesPath, str], BinaryIO] - Filesystem opener, by default `builtins.open` Returns ------- bool Whether the can be opened. """ - if hasattr(path, "read"): - path = cast("BinaryIO", path) + if is_read_seek_binary(path): path.seek(0) magic = path.read(4) else: - with open_(path) as fh: + with open(cast("StrOrPath", path), "rb") as fh: magic = fh.read(4) return magic in (NEW_HEADER_MAGIC, OLD_HEADER_MAGIC) diff --git a/tests/test_reader.py b/tests/test_reader.py index f445fde..dc8f2eb 100644 --- a/tests/test_reader.py +++ b/tests/test_reader.py @@ -1,3 +1,4 @@ +import io import json import pickle import sys @@ -8,6 +9,7 @@ import pytest from nd2 import ND2File, imread from nd2._parse._chunk_decode import get_version +from nd2._readers import protocol from nd2._util import AXIS, is_supported_file from resource_backed_dask_array import ResourceBackedDaskArray @@ -287,3 +289,18 @@ def test_file_handles(single_nd2: Path) -> None: assert isinstance(f.asarray(), np.ndarray) assert fh.closed assert f.closed + + +def test_url_handle_closed_on_reader_create_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Handle(io.BytesIO): + pass + + handle = Handle(b"NOT!") + monkeypatch.setattr(protocol, "open_fsspec_url", lambda *args, **kwargs: handle) + + with pytest.raises(OSError, match="not recognized as ND2"): + protocol.ND2Reader.create("s3://bucket/not-nd2") + + assert handle.closed diff --git a/tests/test_s3.py b/tests/test_s3.py new file mode 100644 index 0000000..14d4b9f --- /dev/null +++ b/tests/test_s3.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +import numpy as np +import pytest +from nd2 import ND2File, imread + +if TYPE_CHECKING: + from collections.abc import Iterable + from pathlib import Path + + from mypy_boto3_s3.service_resource import Bucket + + +@pytest.fixture(scope="session") +def s3_endpoint() -> Iterable[str]: + if TYPE_CHECKING: + from moto.server import ThreadedMotoServer + else: + ThreadedMotoServer = pytest.importorskip("moto.server").ThreadedMotoServer + + # ThreadedMotoServer is a real HTTP server and does not patch boto3's + # credential chain. Provide dummy credentials so boto3 and s3fs don't + # raise NoCredentialsError in environments without ~/.aws/credentials. + os.environ.setdefault("AWS_ACCESS_KEY_ID", "testing") + os.environ.setdefault("AWS_SECRET_ACCESS_KEY", "testing") + os.environ.setdefault("AWS_SECURITY_TOKEN", "testing") + os.environ.setdefault("AWS_SESSION_TOKEN", "testing") + os.environ.setdefault("AWS_DEFAULT_REGION", "us-east-1") + + server = ThreadedMotoServer(port=0) + server.start() + _, port = server.get_host_and_port() + try: + yield f"http://127.0.0.1:{port}" + finally: + server.stop() + + +@pytest.fixture(scope="session") +def s3_bucket(s3_endpoint: str) -> Bucket: + if TYPE_CHECKING: + import boto3 + else: + boto3 = pytest.importorskip("boto3") + s3 = boto3.resource("s3", endpoint_url=s3_endpoint, region_name="us-east-1") + return s3.create_bucket(Bucket="nd2-test") + + +@pytest.fixture() +def any_s3_nd2_url(small_nd2s: Path, s3_bucket: Bucket) -> Iterable[tuple[str, str]]: + """Yield an S3 URL pointing to any nd2, along with the endpoint URL.""" + s3_bucket.upload_file(str(small_nd2s), small_nd2s.name) + endpoint = s3_bucket.meta.client.meta.endpoint_url + yield f"s3://{s3_bucket.name}/{small_nd2s.name}", endpoint + + +@pytest.fixture() +def single_s3_nd2_url(single_nd2: Path, s3_bucket: Bucket) -> Iterable[tuple[str, str]]: + """Yield an S3 URL pointing to `single_nd2`, along with the endpoint URL.""" + s3_bucket.upload_file(str(single_nd2), single_nd2.name) + endpoint = s3_bucket.meta.client.meta.endpoint_url + yield f"s3://{s3_bucket.name}/{single_nd2.name}", endpoint + + +def test_nd2file_reads_from_s3_url( + any_s3_nd2_url: tuple[str, str], small_nd2s: Path +) -> None: + url, endpoint = any_s3_nd2_url + storage_options = { + "client_kwargs": {"endpoint_url": endpoint}, + # necessary for moto's S3 implementation + # to avoid checksum validation errors on multipart uploads + "config_kwargs": {"response_checksum_validation": "when_required"}, + } + remote_nd = ND2File(url, storage_options=storage_options) + local_nd = ND2File(small_nd2s) + with remote_nd, local_nd: + assert remote_nd.path == url + assert remote_nd.shape == local_nd.shape + assert remote_nd.metadata == local_nd.metadata + assert remote_nd.attributes == local_nd.attributes + assert remote_nd.text_info == local_nd.text_info + assert remote_nd.experiment == local_nd.experiment + assert remote_nd.frame_metadata(0) == local_nd.frame_metadata(0) + assert remote_nd.events() == local_nd.events() + assert remote_nd.ome_metadata() == local_nd.ome_metadata() + np.testing.assert_array_equal(remote_nd.read_frame(0), local_nd.read_frame(0)) + if local_nd.binary_data is not None: + assert remote_nd.binary_data is not None + for rb, lb in zip(remote_nd.binary_data, local_nd.binary_data): + np.testing.assert_array_equal(rb.asarray(), lb.asarray()) + for rr, lr in zip(remote_nd.rois.items(), local_nd.rois.items()): + assert rr == lr + + full_local_read = imread(small_nd2s) + full_remote_read = imread(url, storage_options=storage_options) + np.testing.assert_array_equal(full_remote_read, full_local_read) + + +def test_nd2file_reads_from_fsspec_obj( + single_s3_nd2_url: tuple[str, str], single_nd2: Path +) -> None: + """Test that the user can construct their own fsspec file-like for ND2File.""" + if TYPE_CHECKING: + import fsspec + else: + fsspec = pytest.importorskip("fsspec") + + url, endpoint = single_s3_nd2_url + storage_options = { + "client_kwargs": {"endpoint_url": endpoint}, + "config_kwargs": {"response_checksum_validation": "when_required"}, + } + fs = fsspec.filesystem("s3", **storage_options) + with fs.open(url, "rb") as fs_fh: + remote_nd = ND2File(fs_fh, storage_options=storage_options) + local_nd = ND2File(single_nd2) + with remote_nd, local_nd: + assert remote_nd.path == url + assert remote_nd.shape == local_nd.shape + assert remote_nd.metadata == local_nd.metadata + assert remote_nd.attributes == local_nd.attributes + assert remote_nd.text_info == local_nd.text_info + assert remote_nd.experiment == local_nd.experiment + assert remote_nd.frame_metadata(0) == local_nd.frame_metadata(0) + assert remote_nd.events() == local_nd.events() + assert remote_nd.ome_metadata() == local_nd.ome_metadata() + np.testing.assert_array_equal( + remote_nd.read_frame(0), local_nd.read_frame(0) + )