numpy_array_ext.py

from __future__ import annotations
import numpy as np

'''
Known caveats:
 - Calling any numpy function that doesn't modify the array will still trigger the _changed flag to True
'''

class DummyChunk:
    def __init__(self):
        self._changed = False
        
class AmuletArrayWrapper:
    __slots__ = ("_ndarray", "_chunk")

    def __init__(self, _ndarray, chunk: DummyChunk):
        self._ndarray = _ndarray
        self._chunk: DummyChunk = chunk

    def __repr__(self):
        return repr(self._ndarray)

    def __dir__(self):
        return dir(self._ndarray)

    def __eq__(self, other):
        return self._ndarray == other

    def __lt__(self, other):
        return self._ndarray < other

    def __gt__(self, other):
        return self._ndarray > other

    def __le__(self, other):
        return self._ndarray <= other

    def __ge__(self, other):
        return self._ndarray >= other

    def __contains__(self, item):
        return item in self._ndarray

    def __getattr__(self, item):
        result = getattr(self._ndarray, item, None)
        if result is None:
            raise AttributeError
        if callable(result):
            self._chunk._changed = True
        return result

    def __getitem__(self, item):
        return self._ndarray[item]

    def __setitem__(self, key, value):
        self._ndarray[key] = value
        self._chunk._changed = True

    def __setslice__(self, i, j, sequence):
        return self._ndarray[i:j:sequence]


if __name__ == '__main__':
    chunk = DummyChunk()
    arr = AmuletArrayWrapper(np.ndarray((2,3), buffer=np.array([1,2,3,1,2,3]), dtype=int), chunk)
    print(arr[arr >= 2])
    print('arr:',arr, chunk._changed)
    print('arr[arr == 2]:', arr[arr == 2], chunk._changed)
    arr[np.equal(arr, 2)] = 4
    print('arr:', arr, chunk._changed)
    print('arr[arr == 2]:', arr[arr == 2], chunk._changed)
    print('arr[1]:', arr[1], chunk._changed)
    chunk._changed = False
    print('arr:', arr, chunk._changed)
    arr.fill(7)
    print('arr:', arr, chunk._changed)
    chunk._changed = False
    print('arr:', arr, chunk._changed)
    raveled = arr.ravel()
    print('arr:', arr, chunk._changed)
    print('raveled:', raveled)