Merge pull request #461 from Blosc/array-api

Array api

Author: lshaw8317

README_DEVELOPERS.md

@@ -70,3 +70,13 @@ We are using Sphinx for documentation.  You can build the documentation by execu
 [You may need to install the `pandoc` package first: https://pandoc.org/installing.html]
 
 You will find the documentation in the `../html` directory.
+
+## Array API tests compatibility
+
+You can test array API compatibility with the `array-api-tests` module.
+Use the `tests/array-api-xfails.txt` to skip the tests that are not supported
+and run pytest from the `array-api-tests` source dir like this:
+
+``` bash
+ARRAY_API_TESTS_MODULE=blosc2 pytest array_api_tests --xfails-file ${BLOSC2_DIR}/tests/array-api-xfails.txt -xs
+```

pyproject.toml

@@ -37,11 +37,13 @@ dependencies = [
     "ndindex",
     "msgpack",
     "platformdirs",
-    "numexpr; platform_machine != 'wasm32'",
+    "numexpr>=2.12.1; platform_machine != 'wasm32'",
     "py-cpuinfo; platform_machine != 'wasm32'",
     "requests",
 ]
 version = "3.7.3.dev0"
+[project.entry-points."array_api"]
+blosc2 = "blosc2"
 
 
 [project.optional-dependencies]
@@ -99,8 +101,6 @@ exclude = ["bench*", ".github*"]
 
 [tool.ruff]
 line-length = 109
-# Raise complexity from the default 10 to 13
-complexity = 13
 extend-exclude = ["bench"]
 
 [tool.ruff.lint]
@@ -116,7 +116,7 @@ extend-select = [
     "SIM",
     "TC",
     "UP",
-]
+    "C901"]  # enable complexity rule
 ignore = [
     "B028",
     "PT011",
@@ -130,3 +130,7 @@ ignore = [
 
 [tool.ruff.lint.extend-per-file-ignores]
 "tests/**" = ["F841"]
+
+[tool.ruff.lint.mccabe]
+# Raise complexity from the default 10 to 13
+max-complexity = 13

src/blosc2/__init__.py

@@ -13,6 +13,8 @@
 import platform
 from enum import Enum
 
+import numpy as np
+
 # Do the platform check once at module level
 IS_WASM = platform.machine() == "wasm32"
 # IS_WASM = True  # for testing (comment this line out for production)
@@ -23,9 +25,10 @@
 if not IS_WASM:
     import numexpr
 
-from .version import __version__
+from .version import __array_api_version__, __version__
 
 __version__ = __version__
+__array_api_version__ = __array_api_version__
 """
 Python-Blosc2 version.
 """
@@ -164,6 +167,94 @@ class Tuner(Enum):
 """
 The C-Blosc2 version's string."""
 
+
+# For array-api compatibility
+iinfo = np.iinfo
+finfo = np.finfo
+
+# dtypes for array-api
+str_ = np.str_
+bytes_ = np.bytes_
+object_ = np.object_
+
+from numpy import (
+    bool_,
+    complex64,
+    complex128,
+    e,
+    euler_gamma,
+    float16,
+    float32,
+    float64,
+    inf,
+    int8,
+    int16,
+    int32,
+    int64,
+    nan,
+    newaxis,
+    pi,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+)
+
+DEFAULT_COMPLEX = complex128
+"""
+Default complex floating dtype."""
+
+DEFAULT_FLOAT = float64
+"""
+Default real floating dtype."""
+
+DEFAULT_INT = int64
+"""
+Default integer dtype."""
+
+DEFAULT_INDEX = int64
+"""
+Default indexing dtype."""
+
+
+class Info:
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+
+def __array_namespace_info__() -> Info:
+    """
+    Return information about the array namespace following the Array API specification.
+    """
+
+    def _raise(exc):
+        raise exc
+
+    return Info(
+        capabilities=lambda: {
+            "boolean indexing": True,
+            "data-dependent shapes": False,
+            "max dimensions": MAX_DIM,
+        },
+        default_device=lambda: "cpu",
+        default_dtypes=lambda device=None: {
+            "real floating": DEFAULT_FLOAT,
+            "complex floating": DEFAULT_COMPLEX,
+            "integral": DEFAULT_INT,
+            "indexing": DEFAULT_INDEX,
+        }
+        if (device == "cpu" or device is None)
+        else _raise(ValueError("Only cpu devices allowed")),
+        dtypes=lambda device=None, kind=None: np.__array_namespace_info__().dtypes(kind=kind, device=device)
+        if (device == "cpu" or device is None)
+        else _raise(ValueError("Only cpu devices allowed")),
+        devices=lambda: ["cpu"],
+        name="blosc2",
+        version=__version__,
+    )
+
+
 # Public API for container module
 from .core import (
     clib_info,
@@ -237,9 +328,11 @@ class Tuner(Enum):
     are_partitions_aligned,
     are_partitions_behaved,
     arange,
+    broadcast_to,
     linspace,
     eye,
     asarray,
+    astype,
     indices,
     sort,
     reshape,
@@ -248,14 +341,19 @@ class Tuner(Enum):
     concatenate,
     expand_dims,
     empty,
+    empty_like,
     frombuffer,
     fromiter,
     get_slice_nchunks,
+    meshgrid,
     nans,
     uninit,
     zeros,
+    zeros_like,
     ones,
+    ones_like,
     full,
+    full_like,
     save,
     matmul,
     permute_dims,
@@ -278,6 +376,7 @@ class Tuner(Enum):
     get_expr_operands,
     validate_expr,
     evaluate,
+    slices_eval,
 )
 from .proxy import Proxy, ProxySource, ProxyNDSource, ProxyNDField, SimpleProxy, jit
 
@@ -333,9 +432,13 @@ class Tuner(Enum):
     contains,
     cos,
     cosh,
+    equal,
     exp,
     expm1,
     imag,
+    isfinite,
+    isinf,
+    isnan,
     lazywhere,
     log,
     log1p,
@@ -356,14 +459,25 @@ class Tuner(Enum):
     where,
 )
 
-__all__ = [
+__all__ = [  # noqa : RUF022
     # Constants
     "EXTENDED_HEADER_LENGTH",
     "MAX_BUFFERSIZE",
     "MAX_TYPESIZE",
     "MIN_HEADER_LENGTH",
     "VERSION_DATE",
     "VERSION_STRING",
+    # Default dtypes
+    "DEFAULT_COMPLEX",
+    "DEFAULT_FLOAT",
+    "DEFAULT_INDEX",
+    "DEFAULT_INT",
+    # Mathematical constants
+    "e",
+    "pi",
+    "inf",
+    "nan",
+    "newaxis",
     # Classes
     "C2Array",
     "CParams",
@@ -407,7 +521,8 @@ class Tuner(Enum):
     "are_partitions_aligned",
     "are_partitions_behaved",
     "asarray",
-    "clib_info",
+    "astypeclib_info",
+    "broadcast_to",
     "compress",
     "compress2",
     "compressor_list",
@@ -421,6 +536,9 @@ class Tuner(Enum):
     "decompress2",
     "detect_number_of_cores",
     "dparams_dflts",
+    "empty",
+    "empty_like",
+    "equal",
     "estore_from_cframe",
     "expand_dims",
     "expm1",
@@ -430,6 +548,7 @@ class Tuner(Enum):
     "frombuffer",
     "fromiter",
     "full",
+    "full_like",
     "get_blocksize",
     "get_cbuffer_sizes",
     "get_clib",
@@ -438,6 +557,9 @@ class Tuner(Enum):
     "get_expr_operands",
     "get_slice_nchunks",
     "indices",
+    "isfinite",
+    "isinf",
+    "isnan",
     "jit",
     "lazyexpr",
     "lazyudf",
@@ -452,10 +574,12 @@ class Tuner(Enum):
     "matrix_transpose",
     "max",
     "mean",
+    "meshgrid",
     "min",
     "nans",
     "ndarray_from_cframe",
     "ones",
+    "ones_like",
     "open",
     "pack",
     "pack_array",
@@ -494,4 +618,6 @@ class Tuner(Enum):
     "validate_expr",
     "var",
     "where",
+    "zeros",
+    "zeros_like",
 ]

src/blosc2/blosc2_ext.pyx

@@ -2612,8 +2612,8 @@ cdef class NDArray:
                 mask_[i] = mask[i]
             _check_rc(b2nd_squeeze_index(self.array, mask_), "Error while squeezing array")
 
-        if self.array.shape[0] == 1 and self.ndim == 1:
-            self.array.ndim = 0
+        #if self.array.shape[0] == 1 and self.ndim == 1:
+        #    self.array.ndim = 0
 
     def as_ffi_ptr(self):
         return PyCapsule_New(self.array, <char *> "b2nd_array_t*", NULL)

src/blosc2/core.py

@@ -1467,7 +1467,7 @@ def compute_chunks_blocks(  # noqa: C901
 
     # Return an arbitrary value for chunks and blocks when shape has any 0 dim
     if 0 in shape:
-        return (1,) * len(shape), (1,) * len(shape)
+        return shape, shape
 
     if blocks:
         if not isinstance(blocks, tuple | list):