numpy · mattip · May 3, 2018 · May 3, 2018 · May 8, 2018 · May 8, 2018
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
@@ -6277,6 +6277,27 @@ def luf(lamdaexpr, *args, **kwargs):
 
     """))
 
+
+##############################################################################
+#
+# Documentation for ufunc_wrapper
+#
+##############################################################################
+
+add_newdoc('numpy.core', 'ufunc_wraper',
+    """
+    Decorator class to allow a ufunc-like class method to use the `__array_ufunc__`
+    mechanism
+
+    Examples
+    --------
+
+    >>> class ArrayLike:
+    ...     pass
+
+
+
+    """)
 ##############################################################################
 #
 # Documentation for dtype attributes and methods

diff --git a/numpy/core/include/numpy/ufuncobject.h b/numpy/core/include/numpy/ufuncobject.h
@@ -111,15 +111,20 @@ typedef int (PyUFunc_MaskedInnerLoopSelectionFunc)(
                             NpyAuxData **out_innerloopdata,
                             int *out_needs_api);
 
-typedef struct _tagPyUFuncObject {
-        PyObject_HEAD
-        /*
+/* used in both ufunc and ufunc_wrapper */
+#define UFUNC_BASE \
+    PyObject_HEAD /*
          * nin: Number of inputs
          * nout: Number of outputs
          * nargs: Always nin + nout (Why is it stored?)
-         */
-        int nin, nout, nargs;
+ */ int nin, nout, nargs;
 
+typedef struct {
+        UFUNC_BASE
+} PyUFuncBaseObject;
+
+typedef struct _tagPyUFuncObject {
+        UFUNC_BASE
         /* Identity for reduction, either PyUFunc_One or PyUFunc_Zero */
         int identity;
 

diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
@@ -28,14 +28,16 @@
     from .multiarray import newbuffer, getbuffer
 
 from . import umath
-from .umath import (multiply, invert, sin, UFUNC_BUFSIZE_DEFAULT,
+from .umath import (multiply, invert, sin, ufunc_wrapper,
+                    UFUNC_BUFSIZE_DEFAULT,
                     ERR_IGNORE, ERR_WARN, ERR_RAISE, ERR_CALL, ERR_PRINT,
                     ERR_LOG, ERR_DEFAULT, PINF, NAN)
 from . import numerictypes
 from .numerictypes import longlong, intc, int_, float_, complex_, bool_
 from ._internal import TooHardError, AxisError
 
 bitwise_not = invert
+# TODO properly export ufunc from umath
 ufunc = type(sin)
 newaxis = None
 
@@ -48,6 +50,32 @@
     import __builtin__ as builtins
 
 
+# TODO: rewrite the __call__ in C?
+class UFuncWrapper(ufunc_wrapper):
+    ''' A decorator to wrap ufunc-like class methods and enable the
+        __array_ufunc__ protocol
+
+        Use as
+
+        class MyArray(ndarray):
+            @UFuncWrapper(2, 1) # nin, nout
+            def solve(self, other):
+                ...
+                return result
+    '''
+    def __call__(self, meth):
+        def wrap(*args, **kwds):
+            (status, r) = self.check_override(*args, **kwds)
+            if status > 0:
+                return r
+            return meth(*args, **kwds)
+        wrap.__name__ = meth.__name__
+        wrap.__doc__ = meth.__doc__
+        self.__name__ = meth.__name__
+        self.__doc__ = meth.__doc__
+        return wrap
+
+
 def loads(*args, **kwargs):
     # NumPy 1.15.0, 2017-12-10
     warnings.warn(
@@ -74,7 +102,7 @@ def loads(*args, **kwargs):
     'False_', 'True_', 'bitwise_not', 'CLIP', 'RAISE', 'WRAP', 'MAXDIMS',
     'BUFSIZE', 'ALLOW_THREADS', 'ComplexWarning', 'full', 'full_like',
     'matmul', 'shares_memory', 'may_share_memory', 'MAY_SHARE_BOUNDS',
-    'MAY_SHARE_EXACT', 'TooHardError', 'AxisError']
+    'MAY_SHARE_EXACT', 'TooHardError', 'AxisError', 'UFuncWrapper']
 
 if sys.version_info[0] < 3:
     __all__.extend(['getbuffer', 'newbuffer'])

diff --git a/numpy/core/src/multiarray/number.c b/numpy/core/src/multiarray/number.c
@@ -369,13 +369,44 @@ array_divmod(PyArrayObject *m1, PyObject *m2)
 static PyObject *
 array_matrix_multiply(PyArrayObject *m1, PyObject *m2)
 {
-    static PyObject *matmul = NULL;
-
-    npy_cache_import("numpy.core.multiarray", "matmul", &matmul);
-    if (matmul == NULL) {
-        return NULL;
+    static PyObject *matmul=NULL, *wrapper=NULL, *ufunc=NULL, *checker=NULL;
+    PyObject *result = Py_None, *res_tuple;
+    int status;
+    if (ufunc == NULL) {
+        PyObject *s;
+        npy_cache_import("numpy.core.multiarray", "matmul", &matmul);
+        if (matmul == NULL) {
+            return NULL;
+        }
+        npy_cache_import("numpy.core.umath", "ufunc_wrapper", &wrapper);
+        if (wrapper == NULL) {
+            return NULL;
+        }
+        checker = PyObject_GetAttrString(wrapper, "check_override");
+        if (checker == NULL) {
+            return NULL;
+        }
+        s = Py_BuildValue("ii", 2, 1);
+        ufunc = PyObject_CallObject(wrapper, s);
+        Py_DECREF(s);
+        if (ufunc == NULL) {
+            return NULL;
+        }
     }
     BINOP_GIVE_UP_IF_NEEDED(m1, m2, nb_matrix_multiply, array_matrix_multiply);
+    res_tuple = PyObject_CallFunctionObjArgs(checker, ufunc, m1, m2, NULL);
+    if (PyArg_ParseTuple(res_tuple, "iO", &status, &result) < 0) {
+        Py_DECREF(res_tuple);
+        return NULL;
+    }
+    if (status > 0) {
+        Py_INCREF(result);
+        Py_DECREF(res_tuple);
+        return result;
+    }
+    Py_DECREF(res_tuple);
+    if (PyErr_Occurred())
+        return NULL;
     return PyArray_GenericBinaryFunction(m1, m2, matmul);
 }
 

diff --git a/numpy/core/src/private/ufunc_override.c b/numpy/core/src/private/ufunc_override.c
@@ -41,6 +41,12 @@ get_non_default_array_ufunc(PyObject *obj)
         return NULL;
     }
     /* does the class define __array_ufunc__? */
+#if PY_VERSION_HEX < 0x03000000
+    if (Py_TYPE(obj) == &PyInstance_Type) {
+       PyErr_Warn(PyExc_RuntimeWarning,
+                   "cannot lookup __array_ufunc__ on old-style classes");
+    }
+#endif
     cls_array_ufunc = PyArray_LookupSpecial(obj, "__array_ufunc__");
     if (cls_array_ufunc == NULL) {
         return NULL;

diff --git a/numpy/core/src/umath/override.c b/numpy/core/src/umath/override.c
@@ -37,7 +37,7 @@ normalize_signature_keyword(PyObject *normal_kwds)
 }
 
 static int
-normalize___call___args(PyUFuncObject *ufunc, PyObject *args,
+normalize___call___args(PyUFuncBaseObject *ufunc, PyObject *args,
                         PyObject **normal_args, PyObject **normal_kwds)
 {
     /*
@@ -114,7 +114,7 @@ normalize___call___args(PyUFuncObject *ufunc, PyObject *args,
 }
 
 static int
-normalize_reduce_args(PyUFuncObject *ufunc, PyObject *args,
+normalize_reduce_args(PyUFuncBaseObject *ufunc, PyObject *args,
                       PyObject **normal_args, PyObject **normal_kwds)
 {
     /*
@@ -169,7 +169,7 @@ normalize_reduce_args(PyUFuncObject *ufunc, PyObject *args,
 }
 
 static int
-normalize_accumulate_args(PyUFuncObject *ufunc, PyObject *args,
+normalize_accumulate_args(PyUFuncBaseObject *ufunc, PyObject *args,
                           PyObject **normal_args, PyObject **normal_kwds)
 {
     /*
@@ -215,7 +215,7 @@ normalize_accumulate_args(PyUFuncObject *ufunc, PyObject *args,
 }
 
 static int
-normalize_reduceat_args(PyUFuncObject *ufunc, PyObject *args,
+normalize_reduceat_args(PyUFuncBaseObject *ufunc, PyObject *args,
                     PyObject **normal_args, PyObject **normal_kwds)
 {
     /*
@@ -263,7 +263,7 @@ normalize_reduceat_args(PyUFuncObject *ufunc, PyObject *args,
 }
 
 static int
-normalize_outer_args(PyUFuncObject *ufunc, PyObject *args,
+normalize_outer_args(PyUFuncBaseObject *ufunc, PyObject *args,
                      PyObject **normal_args, PyObject **normal_kwds)
 {
     /*
@@ -297,7 +297,7 @@ normalize_outer_args(PyUFuncObject *ufunc, PyObject *args,
 }
 
 static int
-normalize_at_args(PyUFuncObject *ufunc, PyObject *args,
+normalize_at_args(PyUFuncBaseObject *ufunc, PyObject *args,
                   PyObject **normal_args, PyObject **normal_kwds)
 {
     /* ufunc.at(a, indices[, b]) */
@@ -325,7 +325,7 @@ normalize_at_args(PyUFuncObject *ufunc, PyObject *args,
  * result of the operation, if any. If *result is NULL, there is no override.
  */
 NPY_NO_EXPORT int
-PyUFunc_CheckOverride(PyUFuncObject *ufunc, char *method,
+PyUFunc_CheckOverride(PyUFuncBaseObject *ufunc, char *method,
                       PyObject *args, PyObject *kwds,
                       PyObject **result)
 {

diff --git a/numpy/core/src/umath/override.h b/numpy/core/src/umath/override.h
@@ -5,7 +5,7 @@
 #include "numpy/ufuncobject.h"
 
 NPY_NO_EXPORT int
-PyUFunc_CheckOverride(PyUFuncObject *ufunc, char *method,
+PyUFunc_CheckOverride(PyUFuncBaseObject *ufunc, char *method,
                       PyObject *args, PyObject *kwds,
                       PyObject **result);
 #endif