CODEWITHSUNDEEP
pythonnumpysubclassing
Ben Whale
Ben Whale

Reputation: 493

How to subclass a subclass of numpy.ndarray

I'm struggling to subclass my own subclass of numpy.ndarray. I don't really understand what the problem is and would like someone to explain what goes wrong in the following cases and how to do what I'm trying to do.

What I'm trying to achieve:

I have a subclass of numpy.ndarry that behaves as I want (class A in the code below). I want to subclass A (class B in the code below) so that B contains additional information (name) and methods (the decorated .simple_data method).

Case 1:

import numpy as np

class A(np.ndarray):

    def __new__(cls,data):
        obj = np.asarray(data).view(cls)
        return obj

    def __array_finalize(self,obj):
        if obj is None: return

class B(A):

    def __init__(self,data,name):
        super(B,self).__init__(data)
        self.name = name

    @property
    def simple_data(self):
        return [data[0,:],data[:,0]]

if __name__ == '__main__':
    data = np.arange(20).reshape((4,5))
    b = B(data,'B')
    print type(b)
    print b.simple_data

Running this code produces the output:

Traceback (most recent call last):
  File "ndsubclass.py", line 24, in <module>
    b = B(data,'B')
TypeError: __new__() takes exactly 2 arguments (3 given)

I assume that this is related to the 'name' variable in the construction of B and that due to A being a subclass of numpy.array, A's new method is being called before B's init method. Thus to fix this I assume that B also needs a new method that appropriately handles the additional argument.

My guess is something like:

def __new__(cls,data,name):
    obj = A(data)
    obj.name = name
    return obj

should do it, but how do I change the class of obj?

Case 2:

import numpy as np

class A(np.ndarray):

    def __new__(cls,data):
        obj = np.asarray(data).view(cls)
        return obj

    def __array_finalize__(self,obj):
        if obj is None: return

class B(A):

    def __new__(cls,data):
        obj = A(data)
        obj.view(cls)
        return obj

    def __array_finalize__(self,obj):
        if obj is None: return

    @property
    def simple_data(self):
        return [self[0,:],self[:,0]]

if __name__ == '__main__':
    data = np.arange(20).reshape((4,5))
    b = B(data)
    print type(b)
    print b.simple_data()

When run the output is:

<class '__main__.A'>
Traceback (most recent call last):
  File "ndsubclass.py", line 30, in <module>
    print b.simple_data()
AttributeError: 'A' object has no attribute 'simple_data'

This surprises me as I was expecting:

<class '__main__.B'>
[array([0, 1, 2, 3, 4]), array([ 0,  5, 10, 15])]

I assume that the call to view() in B.new() is somehow not correctly setting the class of obj. Why?

I'm confused as to what is going on and would be very grateful if someone could explain it.

Upvotes: 6

Views: 2330

Answers (1)

agf
agf

Reputation: 176960

For Case 1, the simplest way is:

class B(A):
    def __new__(cls,data,name):
        obj = A.__new__(cls, data)
        obj.name = name
        return obj

__new__ is actually a static method that takes a class as the first argument, not a class method, so you can call it directly with the class of which you want to create an instance.

For Case 2, view doesn't work in-place, you need to assign the result to something, the simplest way is:

class B(A):
    def __new__(cls,data):
        obj = A(data)
        return obj.view(cls)

Also, you've got __array_finalize__ defined the same in A and B there (probably just a typo) -- you don't need to do that.

Upvotes: 4

Related Questions