6

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.

Ben Whale
  • 493
  • 2
  • 9

1 Answers1

4

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.

agf
  • 171,228
  • 44
  • 289
  • 238
  • Ah... So I can force A.__new__ to view cast the array data into the correct class by passing that class to A.__new__. Is this the pythonic way to do this? – Ben Whale Sep 08 '11 at 03:50
  • 1
    @Ben I would use `super(B, cls).__new__` instead of `A.__new__`, but yes, this is fairly standard. You could also change `A` so it could handle a variable number of arguments or something similar, but that might be outside the scope of this question. – agf Sep 08 '11 at 03:56
  • Great! So are my guesses correct for the two cases mentioned above? – Ben Whale Sep 08 '11 at 04:26
  • Ok... got it. Thanks very much for all the help! If I could up vote your answer I would. – Ben Whale Sep 08 '11 at 05:52