"""

The array class for sharpy, a distributed implementation of the

array API as defined here: https://data-apis.org/array-api/latest

"""

#

# See __init__.py for an implementation overview

#

from . import _sharpy as _csp

from . import array_api as api

def slicefy(x):

if isinstance(x, slice):

return x

# slice that extracts a single element at index x

next_val = None if x == -1 else x + 1

return slice(x, next_val, 1)

class ndarray:

def __init__(self, t):

self._t = t

def __repr__(self):

return self._t.__repr__()

for method in api.api_categories["EWBinOp"]:

if method.startswith("__"):

METHOD = method.upper()

exec(

f"{method} = lambda self, other: ndarray(_csp.EWBinOp.op(_csp.{METHOD}, self._t, other._t if isinstance(other, ndarray) else other))"

)

for method in api.api_categories["EWUnyOp"]:

if method.startswith("__"):

METHOD = method.upper()

exec(

f"{method} = lambda self: ndarray(_csp.EWUnyOp.op(_csp.{METHOD}, self._t))"

)

for method in api.api_categories["UnyOp"]:

exec(f"{method} = lambda self: self._t.{method}()")

for att in api.attributes:

exec(f"{att} = property(lambda self: self._t.{att})")

def astype(self, dtype, copy=False):

return ndarray(self._t.astype(dtype, copy))

def to_device(self, device=""):

return ndarray(self._t.to_device(device))

def __getitem__(self, key):

key = key if isinstance(key, tuple) else (key,)

key = [slicefy(x) for x in key]

return ndarray(self._t.__getitem__(key))

def __setitem__(self, key, value):

key = key if isinstance(key, tuple) else (key,)

key = [slicefy(x) for x in key]

if isinstance(value, ndarray) and value._t.dtype != self._t.dtype:

raise ValueError(

f"Mismatching data type in setitem: {value._t.dtype}, expecting {self._t.dtype}"

)

self._t.__setitem__(

key, value._t if isinstance(value, ndarray) else value

)