[Numpy-discussion] Problem with _wrapit

[Tim Hochberg]

There's a problem with the _wrapit function. If the function that it 
wraps, returns  a scalar, it blows up. Here's an example:

    import numpy

    class narray(numpy.ndarray):
        def get_argmax(self):
            raise AttributeError
        argmax = property(get_argmax)
     
       
    a = narray([3], int, numpy.arange(3))
    print type(a), isinstance(a, numpy.ndarray)

    print a
    print numpy.argmax(a)

==>

    <class '__main__.narray'> True
    [0 1 2]
    Traceback (most recent call last):
      File "find_bug.py", line 13, in ?
        print numpy.argmax(a)
      File "C:\python24\lib\site-packages\numpy\core\oldnumeric.py",
    line 278, in argmax
        return _wrapit(a, 'argmax', axis)
      File "C:\python24\lib\site-packages\numpy\core\oldnumeric.py",
    line 170, in _wrapit
        result = wrap(result)
    TypeError: can only be called with ndarray object


A possible fix is to do another isinstance check on the way out:

    def _wrapit(obj, method, *args, **kwds):
        try:
            wrap = obj.__array_wrap__
        except AttributeError:
            wrap = None
        result = getattr(asarray(obj),method)(*args, **kwds)
        if wrap and isinstance(result, mu.ndarray):
            result = wrap(result)
        return result

That fixes this problem, and I think it should be OK in general, but I'd 
like a second opinion before comitting it.

regards,

-tim