quickly create ndarray subtypes with attributes

array.py

# author: Nicolas Tessore <[email protected]>
# license: MIT
'''utility functions for numpy arrays'''

import numpy as np


def array_with_attributes(name, **attrs):
    '''create an array subclass with given attributes'''

    def __new__(cls, array):
        obj = np.asarray(array).view(cls)
        for a, v in attrs.items():
            setattr(obj, a, v)
        return obj

    def __array_finalize__(self, obj):
        if obj is None:
            return
        for a, v in attrs.items():
            setattr(self, a, getattr(obj, a, v))

    ns = {
        '__new__': __new__,
        '__array_finalize__': __array_finalize__,
    }

    def makeprop(a):
        _a = f'_{a}'
        return property(lambda self: getattr(self, _a),
                        lambda self, v: setattr(self, _a, v))

    for a in attrs:
        ns[a] = makeprop(a)

    return type(name, (np.ndarray,), ns)

test_array.py

def test_array_with_attributes():
    import numpy as np
    from glass.util.array import array_with_attributes

    TestArray = array_with_attributes('TestArray', foo=1, bar=2)
    TestArray.__doc__ = 'the class docstring'
    TestArray.foo.__doc__ = 'the foo docstring'
    TestArray.bar.__doc__ = 'the bar docstring'

    assert issubclass(TestArray, np.ndarray)
    assert TestArray.__name__ == 'TestArray'
    assert hasattr(TestArray, 'foo')
    assert hasattr(TestArray, 'bar')

    a = TestArray([])

    assert isinstance(a, TestArray)
    assert isinstance(a, np.ndarray)

    assert isinstance(a[0:], TestArray)

    b = np.arange(3).view(TestArray)

    assert isinstance(b, TestArray)

    assert a.foo == 1
    assert a.bar == 2

    b.foo = 3

    assert a.foo == 1
    assert b.foo == 3