diff --git a/pyproject.toml b/pyproject.toml index 2506796..989244e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ dev = [ "build", "mypy", + "numpy", "pytest", "ruff", "toml", diff --git a/src/appose/shm.py b/src/appose/shm.py index 528c4af..296dde3 100644 --- a/src/appose/shm.py +++ b/src/appose/shm.py @@ -8,8 +8,8 @@ from __future__ import annotations -import re -from math import ceil, prod +import warnings +from math import prod from multiprocessing import resource_tracker, shared_memory from typing import TYPE_CHECKING @@ -129,15 +129,21 @@ def __init__(self, dtype: str, shape: list[int], shm: SharedMemory | None = None Args: dtype: The type of the data elements; e.g. int8, uint8, float32, float64. + NumPy-style short forms (e.g. u2, f4, |u1, =c8) are also accepted, + and normalized to the standard name (e.g. uint16, float32). + Explicit byte orders (< or >) are rejected: Appose arrays + always use the machine's native byte order. To match a NumPy + array, pass arr.dtype.name, not str(arr.dtype), which keeps a + non-native byte order; or use copy_of to copy the array. shape: The dimensional extents; e.g. a stack of 7 image planes with resolution 512x512 would have shape [7, 512, 512]. shm: The SharedMemory containing the array data, or None to create it. """ - self.dtype: str = dtype + self.dtype: str = _normalize_dtype(dtype) self.shape: list[int] = shape self.shm: SharedMemory = ( SharedMemory( - create=True, rsize=ceil(prod(shape) * _bytes_per_element(dtype)) + create=True, rsize=prod(shape) * _bytes_per_element(self.dtype) ) if shm is None else shm @@ -151,20 +157,62 @@ def __str__(self): f"shm='{self.shm.name}' ({self.shm.rsize}))" ) - def ndarray(self): + def __array__(self, dtype=None, copy=None): """ - Create a NumPy ndarray object for working with the array data. - No array data is copied; the NumPy array wraps the same SharedMemory. + Support numpy.asarray(nda), which wraps the array data as a NumPy + ndarray without copying it; the NumPy array uses the same SharedMemory. Requires the numpy package to be installed. """ try: import numpy - - return numpy.ndarray( - prod(self.shape), dtype=self.dtype, buffer=self.shm.buf - ).reshape(self.shape) except ModuleNotFoundError: raise ImportError("NumPy is not available.") + arr = numpy.ndarray( + prod(self.shape), dtype=self.dtype, buffer=self.shm.buf + ).reshape(self.shape) + if dtype is None: + dtype = arr.dtype + if copy is False and numpy.dtype(dtype) != arr.dtype: + raise ValueError( + f"Cannot convert NDArray from {arr.dtype} to {dtype} without copying" + ) + return arr.astype(dtype, copy=bool(copy)) + + def ndarray(self): + """ + Create a NumPy ndarray object for working with the array data. + No array data is copied; the NumPy array wraps the same SharedMemory. + Requires the numpy package to be installed. + + Deprecated: use numpy.asarray(nda) instead. + """ + warnings.warn( + "NDArray.ndarray() is deprecated; use numpy.asarray(nda) instead", + DeprecationWarning, + stacklevel=2, + ) + return self.__array__() + + @classmethod + def copy_of(cls, arr) -> NDArray: + """ + Create an NDArray in new shared memory, holding a copy of the given + NumPy array. + + The data is copied value by value, so the source array may be in any + byte order and memory layout (e.g. a big-endian array, or a transposed + view); the copy is always C-ordered, in native byte order. + + Args: + arr: The NumPy array to copy. + """ + nda = cls(arr.dtype.name, list(arr.shape)) + try: + nda.__array__()[:] = arr + except BaseException: + nda.shm.dispose() + raise + return nda def __enter__(self) -> Self: return self @@ -187,9 +235,77 @@ def __exit__(self, exc_type, exc_value, exc_tb) -> None: ) -def _bytes_per_element(dtype: str) -> int | float: - try: - bits = int(re.sub("[^0-9]", "", dtype)) - except ValueError: - raise ValueError(f"Invalid dtype: {dtype}") - return bits / 8 +# Standard dtype names, with the number of bytes per element of each. +_DTYPE_SIZES = { + "int8": 1, + "int16": 2, + "int32": 4, + "int64": 8, + "uint8": 1, + "uint16": 2, + "uint32": 4, + "uint64": 8, + "float16": 2, + "float32": 4, + "float64": 8, + "complex64": 8, + "complex128": 16, + "bool": 1, +} + +# NumPy-style short forms of the standard dtype names. +_DTYPE_ALIASES = { + "i1": "int8", + "i2": "int16", + "i4": "int32", + "i8": "int64", + "u1": "uint8", + "u2": "uint16", + "u4": "uint32", + "u8": "uint64", + "f2": "float16", + "f4": "float32", + "f8": "float64", + "c8": "complex64", + "c16": "complex128", + "b1": "bool", + "?": "bool", +} + + +def _normalize_dtype(dtype: str) -> str: + """ + Return the standard name of the given dtype; e.g. " "uint16". + + Accepts standard names (e.g. uint16, float32) as well as NumPy-style + short forms (e.g. u2, f4). A short form may be prefixed with = (native + byte order) or | (byte order not applicable), which is ignored. Explicit + byte orders (< or >) are rejected, so that parsing behaves the same on + every machine; Appose arrays always use the machine's native byte order. + + Only platform-independent types are supported; e.g. longdouble and + single-character codes like "l" are rejected, since their sizes vary. + """ + if dtype in _DTYPE_SIZES: + return dtype + if dtype.startswith(("<", ">")): + name = _DTYPE_ALIASES.get(dtype[1:], dtype[1:]) + if name in _DTYPE_SIZES: + raise ValueError( + f"Unsupported dtype: {dtype} " + "(Appose arrays are always in native byte order; " + f"use '{name}' instead, e.g. via arr.dtype.name, " + "or copy a NumPy array into shared memory " + "via NDArray.copy_of(arr))" + ) + short = dtype[1:] if dtype.startswith(("=", "|")) else dtype + if short not in _DTYPE_ALIASES: + raise ValueError(f"Unsupported dtype: {dtype}") + return _DTYPE_ALIASES[short] + + +def _bytes_per_element(dtype: str) -> int: + """ + Return the number of bytes per element for the given dtype. + """ + return _DTYPE_SIZES[_normalize_dtype(dtype)] diff --git a/tests/test_service.py b/tests/test_service.py index d4e952f..c9029a0 100644 --- a/tests/test_service.py +++ b/tests/test_service.py @@ -247,7 +247,9 @@ def test_python_sys_exit(): def test_crash_with_active_task(): env = appose.system() - with env.python() as service: + # Note: Import numpy (a dev dependency) up front, so that the worker does + # not emit its numpy warning, which would pollute the stderr under test. + with env.python().init("import numpy") as service: maybe_debug(service) # Create a "long-running" task. script = ( diff --git a/tests/test_shm.py b/tests/test_shm.py index 7c50093..29f34dc 100644 --- a/tests/test_shm.py +++ b/tests/test_shm.py @@ -2,8 +2,12 @@ # Copyright (C) 2023 - 2026 Appose developers. # SPDX-License-Identifier: BSD-2-Clause +import numpy +import pytest + import appose from appose.service import TaskStatus +from appose.shm import _bytes_per_element, _normalize_dtype ndarray_inspect = """ task.outputs["rsize"] = data.shm.rsize @@ -39,3 +43,148 @@ def test_ndarray(): assert "uint16" == task.outputs["dtype"] assert [2, 20, 25] == task.outputs["shape"] assert 123 + 78 + 210 == task.outputs["sum"] + + +def test_dtype_standard_names(): + for dtype, size in [ + ("int8", 1), + ("int16", 2), + ("int32", 4), + ("int64", 8), + ("uint8", 1), + ("uint16", 2), + ("uint32", 4), + ("uint64", 8), + ("float16", 2), + ("float32", 4), + ("float64", 8), + ("complex64", 8), + ("complex128", 16), + ("bool", 1), + ]: + assert dtype == _normalize_dtype(dtype) + assert size == _bytes_per_element(dtype) + + +def test_dtype_short_forms(): + for short, name in [ + ("i1", "int8"), + ("i2", "int16"), + ("i4", "int32"), + ("i8", "int64"), + ("u1", "uint8"), + ("u2", "uint16"), + ("u4", "uint32"), + ("u8", "uint64"), + ("f2", "float16"), + ("f4", "float32"), + ("f8", "float64"), + ("c8", "complex64"), + ("c16", "complex128"), + ("b1", "bool"), + ("?", "bool"), + ]: + assert name == _normalize_dtype(short) + assert name == _normalize_dtype("=" + short) + assert name == _normalize_dtype("|" + short) + + +def test_dtype_explicit_byte_order(): + for dtype, name in [ + ("u2", "uint16"), + ("c16", "complex128"), + ("float32", "float32"), + ]: + with pytest.raises(ValueError, match=f"native byte order; use '{name}'"): + _normalize_dtype(dtype) + + +def test_dtype_unsupported(): + for dtype in [ + "", + "=", + "==u2", + "|=u2", + "=uint16", + "|uint8", + "uint", + "u3", + "f16", + "c32", + "?1", + "l", + "g", + "longdouble", + "float128", + "intp", + "U10", + "i3", + "datetime64[ns]", + "object", + "FLOAT32", + ]: + with pytest.raises(ValueError, match="Unsupported dtype"): + _normalize_dtype(dtype) + + +def test_ndarray_normalizes_dtype(): + with appose.NDArray("=u2", [3, 5]) as data: + assert "uint16" == data.dtype + assert 3 * 5 * 2 == data.shm.rsize + + +def test_copy_of_big_endian(): + # A big-endian array, as produced by some image readers. + src = (numpy.arange(3 * 4 * 5).reshape(3, 4, 5) * 1000).astype(">u2") + with pytest.raises(ValueError, match="use 'uint16'"): + appose.NDArray(str(src.dtype), list(src.shape)) + with appose.NDArray.copy_of(src) as data: + assert "uint16" == data.dtype + assert [3, 4, 5] == data.shape + assert 3 * 4 * 5 * 2 == data.shm.rsize + dst = numpy.asarray(data) + assert dst.dtype.isnative + assert numpy.array_equal(src, dst) + + +def test_copy_of_non_contiguous(): + src = numpy.arange(24, dtype="float32").reshape(2, 3, 4).transpose(2, 0, 1) + with appose.NDArray.copy_of(src) as data: + assert "float32" == data.dtype + assert [4, 2, 3] == data.shape + assert numpy.array_equal(src, numpy.asarray(data)) + + +def test_copy_of_unsupported(): + with pytest.raises(ValueError, match="Unsupported dtype: datetime64"): + appose.NDArray.copy_of(numpy.zeros(3, dtype="datetime64[ns]")) + + +def test_asarray_zero_copy(): + with appose.NDArray("float32", [2, 3]) as data: + arr = numpy.asarray(data) + assert "float32" == arr.dtype.name + assert (2, 3) == arr.shape + arr[1, 2] = 42 + assert 42 == numpy.asarray(data)[1, 2] + + copied = numpy.array(data) + copied[0, 0] = 7 + assert 0 == numpy.asarray(data)[0, 0] + + converted = numpy.asarray(data, dtype="float64") + assert "float64" == converted.dtype.name + assert 42 == converted[1, 2] + + +def test_ndarray_deprecated(): + with appose.NDArray("uint8", [4]) as data: + with pytest.warns(DeprecationWarning, match="numpy.asarray"): + arr = data.ndarray() + arr[0] = 9 + assert 9 == numpy.asarray(data)[0]