Will the Array API promise a uniform way to extend arrays?

[NeilGirdhar]

Problem

I would like to extend an array type to have some user-defined flags.

Background

Various array libraries provide some mechanisms for extending their arrays. For example, NumPy provides __array_function__ described in NEP 18, and Jax provides an undocumented/unrecommended __jax_array__.

Example

Suppose I would like to add flags to an array type. Would something like this work?

class HasFlags:
    def __init__(self, flags: dict[str, Any], **kwargs) -> None:
        super().__init__(**kwargs)
        self.flags = flags


class ArrayAPIBase:   # Perhaps have this in array-api-compat?
    """A base class to implement an Array API array by composition."""
    def __init__(self, array: Array, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self.array = array

    def __array_namespace__(self):  # noqa: ANN204, PLW3201
        return self.array.__array_namespace__

    dtype = property(lambda self: self.array.dtype)
    shape = property(lambda self: self.array.shape)
    # ...


class FlagArray(ArrayAPIBase, HasFlags):
    pass

[rgommers]

I think extending by either subclassing or a protocol a la __array_function__ is impossible to standardize - too hard, no commonality between libraries, and has runtime overhead.

What can be done is wrapping with a new array class, stashing the array API standard-compliant array as a property on a wrapper array instance. @mdhaber has an implementation of masked arrays that do such a thing at https://github.com/mdhaber/marray. I believe it's (almost) fully functional.