Improve registered ufuncs so that they can coerce data to and from user-defined data-types. Clean-up _SOFFSET_ macro. Make sure that user-defined types always succeed in CanCastSafely. This is a hack but there is currently no way to record this information.

5 files changed, 108 insertions, 69 deletions

diff --git a/numpy/core/include/numpy/arrayobject.h b/numpy/core/include/numpy/arrayobject.h
index d337906ca..d713666b9 100644
--- a/numpy/core/include/numpy/arrayobject.h
+++ b/numpy/core/include/numpy/arrayobject.h
@@ -69,7 +69,7 @@ extern "C" CONFUSE_EMACS
 /* allocated statically.  This is the size of that static allocation. */
 /*  The array creation itself could have arbitrary dimensions but
  *  all the places where static allocation is used would need to
- *  be changed to dynamic (including inside of structures)
+ *  be changed to dynamic (including inside of several structures)
  */
 
 #define MAX_DIMS 32
diff --git a/numpy/core/src/arrayobject.c b/numpy/core/src/arrayobject.c
index 9fbffb3c3..5b51b5c15 100644
--- a/numpy/core/src/arrayobject.c
+++ b/numpy/core/src/arrayobject.c
@@ -1118,7 +1118,7 @@ PyArray_Scalar(void *data, PyArray_Descr *descr, PyObject *base)
 		}
 	}
 	else {
-		destptr = _SOFFSET_(obj, type_num);
+		destptr = _SOFFSET_(obj);
 	}
 	/* copyswap for OBJECT increments the reference count */
         copyswap(destptr, data, swap, base);
@@ -7186,6 +7186,12 @@ PyArray_CanCastTo(PyArray_Descr *from, PyArray_Descr *to)
 		    stringified value of the object.
 		*/
 	}
+	if (!ret && (PyTypeNum_ISUSERDEF(fromtype) ||	\
+		     PyTypeNum_ISUSERDEF(totype))) {
+		/* don't restrict casting to and from
+		   additional data-types */
+		return TRUE;
+	}
 	return ret;
 }
 
diff --git a/numpy/core/src/multiarraymodule.c b/numpy/core/src/multiarraymodule.c
index 0ad681e7e..1c10f0c33 100644
--- a/numpy/core/src/multiarraymodule.c
+++ b/numpy/core/src/multiarraymodule.c
@@ -1658,8 +1658,6 @@ PyArray_CanCoerceScalar(char thistype, char neededtype,
 }
 
 
-/* This needs to change to allow scalars of a different "kind" to alter the input type
- */
 
 /*OBJECT_API*/
 static PyArrayObject **
@@ -1747,7 +1745,6 @@ PyArray_ConvertToCommonType(PyObject *op, int *retn)
 
 
 /*MULTIARRAY_API
- Numeric.choose()
 */
 static PyObject *
 PyArray_Choose(PyArrayObject *ip, PyObject *op)
diff --git a/numpy/core/src/scalartypes.inc.src b/numpy/core/src/scalartypes.inc.src
index b4622bb32..ea5feb976 100644
--- a/numpy/core/src/scalartypes.inc.src
+++ b/numpy/core/src/scalartypes.inc.src
@@ -5,7 +5,7 @@
 #endif
 #include "numpy/arrayscalars.h"
 
-#define _SOFFSET_(obj,y) ((char *)obj + (offsetof(PyScalarObject, obval)))
+#define _SOFFSET_(obj) ((char *)obj + (offsetof(PyScalarObject, obval)))
 
 static PyBoolScalarObject _PyArrayScalar_BoolValues[2] = {
 	{PyObject_HEAD_INIT(&PyBoolArrType_Type) 0},
@@ -58,8 +58,7 @@ PyArray_ScalarAsCtype(PyObject *scalar, void *ctypeptr)
 		}
 		return;
 	}
-	memcpy(ctypeptr, _SOFFSET_(scalar, typecode->type_num),
-	       typecode->elsize);
+	memcpy(ctypeptr, _SOFFSET_(scalar), typecode->elsize);
 	Py_DECREF(typecode);
 	return;
 }
@@ -77,8 +76,11 @@ PyArray_CastScalarToCtype(PyObject *scalar, void *ctypeptr,
 			  PyArray_Descr *outcode)
 {
 	PyArray_Descr* descr;
+	PyArray_VectorUnaryFunc* castfunc;
 	
 	descr = PyArray_DescrFromScalar(scalar);
+	castfunc = PyArray_GetCastFunc(descr, outcode->type_num);
+	if (castfunc == NULL) return -1;
 	if (PyTypeNum_ISEXTENDED(descr->type_num) ||
 	    PyTypeNum_ISEXTENDED(outcode->type_num)) {
 		PyArrayObject *ain, *aout;
@@ -92,15 +94,12 @@ PyArray_CastScalarToCtype(PyObject *scalar, void *ctypeptr,
 					     NULL, ctypeptr, 
 					     CARRAY_FLAGS, NULL);
 		if (aout == NULL) {Py_DECREF(ain); return -1;}
-		descr->f->cast[outcode->type_num](ain->data, 
-					       aout->data, 1, ain, aout);
+		castfunc(ain->data, aout->data, 1, ain, aout);
 		Py_DECREF(ain);
 		Py_DECREF(aout);
 	}
 	else {
-		descr->f->cast[outcode->type_num](_SOFFSET_(scalar, 
-							 descr->type_num), 
-					       ctypeptr, 1, NULL, NULL);
+		castfunc(_SOFFSET_(scalar), ctypeptr, 1, NULL, NULL);
 	}
 	Py_DECREF(descr);
 	return 0;
@@ -113,8 +112,10 @@ static int
 PyArray_CastScalarDirect(PyObject *scalar, PyArray_Descr *indescr,
 			 void *ctypeptr, int outtype)
 {
-	indescr->f->cast[outtype](_SOFFSET_(scalar, indescr->type_num),
-				  ctypeptr, 1, NULL, NULL);
+	PyArray_VectorUnaryFunc* castfunc;
+	castfunc = PyArray_GetCastFunc(indescr, outtype);
+	if (castfunc == NULL) return -1;
+	castfunc(_SOFFSET_(scalar), ctypeptr, 1, NULL, NULL);
 	return 0;
 }
 
@@ -177,7 +178,7 @@ PyArray_FromScalar(PyObject *scalar, PyArray_Descr *outcode)
 			memptr = (((PyVoidScalarObject *)scalar)->obval);
 		}
 		else {
-			memptr = _SOFFSET_(scalar, typecode->type_num);
+			memptr = _SOFFSET_(scalar);
 		}
 		break;
 	}
@@ -830,7 +831,7 @@ gentype_real_get(PyObject *self)
 
 	if (PyArray_IsScalar(self, ComplexFloating)) {
 		typecode = _realdescr_fromcomplexscalar(self, &typenum);
-		ret = PyArray_Scalar(_SOFFSET_(self, typenum), typecode,
+		ret = PyArray_Scalar(_SOFFSET_(self), typecode,
 				     NULL);
 		Py_DECREF(typecode);
 		return ret;
@@ -854,7 +855,7 @@ gentype_imag_get(PyObject *self)
 	
 	typecode = _realdescr_fromcomplexscalar(self, &typenum);
 	if (PyArray_IsScalar(self, ComplexFloating)) {
-		ret = PyArray_Scalar(_SOFFSET_(self, typenum)		\
+		ret = PyArray_Scalar(_SOFFSET_(self)		\
 				     + typecode->elsize, typecode, NULL);
 	}
 	else if PyArray_IsScalar(self, Object) {
@@ -1099,8 +1100,7 @@ voidtype_getfield(PyVoidScalarObject *self, PyObject *args, PyObject *kwds)
 		PyArray_Descr *new;
 		if (!PyArray_ISNBO(self->descr->byteorder)) {
 			new = PyArray_DescrFromScalar(ret);
-			byte_swap_vector(_SOFFSET_(ret, new->type_num), 1,
-					 new->elsize);
+			byte_swap_vector(_SOFFSET_(ret), 1, new->elsize);
 			Py_DECREF(new);
 		}
 	}
@@ -1529,7 +1529,7 @@ gentype_getreadbuf(PyObject *self, int segment, void **ptrptr)
 			*ptrptr = ((PyVoidScalarObject *)self)->obval;
 	}
 	else 
-		*ptrptr = (void *)_SOFFSET_(self, outcode->type_num);
+		*ptrptr = (void *)_SOFFSET_(self);
 
 	Py_DECREF(outcode);
 	return numbytes;
diff --git a/numpy/core/src/ufuncobject.c b/numpy/core/src/ufuncobject.c
index 25d84d8fd..a4e284ccc 100644
--- a/numpy/core/src/ufuncobject.c
+++ b/numpy/core/src/ufuncobject.c
@@ -625,52 +625,66 @@ select_types(PyUFuncObject *self, int *arg_types,
              PyUFuncGenericFunction *function, void **data,
 	     PyArray_SCALARKIND *scalars)
 {
-
 	int i=0, j;
 	char start_type;
-	
-	if (PyTypeNum_ISUSERDEF((arg_types[0]))) {
-		PyObject *key, *obj;
+
+	if (self->userloops) {
+		int userdef=-1;
 		for (i=0; i<self->nin; i++) {
-			if (arg_types[i] != arg_types[0]) {
-				PyErr_SetString(PyExc_TypeError,
-						"ufuncs on user defined" \
-						" types don't support "\
-						"coercion");
-				return -1;
+			if (PyTypeNum_ISUSERDEF(arg_types[i])) {
+				userdef = arg_types[i];
+				break;
 			}
 		}
-		for (i=self->nin; i<self->nargs; i++) {
-			arg_types[i] = arg_types[0];
-		}
-
-		obj = NULL;
-		if (self->userloops) {
-			key = PyInt_FromLong((long) arg_types[0]);
+		if (userdef > 0) {
+			PyObject *key, *obj;
+			int *this_types=NULL;
+			
+			obj = NULL;
+			key = PyInt_FromLong((long) userdef);
 			if (key == NULL) return -1;
 			obj = PyDict_GetItem(self->userloops, key);
 			Py_DECREF(key);
+			if (obj == NULL) {
+				PyErr_SetString(PyExc_TypeError, 
+						"user-defined type used in ufunc" \
+						" with no registered loops");
+				return -1;
+			}
+			if PyTuple_Check(obj) {
+				PyObject *item;
+				*function = (PyUFuncGenericFunction)	\
+					PyCObject_AsVoidPtr(PyTuple_GET_ITEM(obj, 
+									     0));
+				item = PyTuple_GET_ITEM(obj, 2);
+				if (PyCObject_Check(item)) {
+					*data = PyCObject_AsVoidPtr(item);
+				}
+				item = PyTuple_GET_ITEM(obj, 1);
+				if (PyCObject_Check(item)) {
+					this_types = PyCObject_AsVoidPtr(item);
+				}
+			}
+			else {
+				*function = (PyUFuncGenericFunction)	\
+					PyCObject_AsVoidPtr(obj);
+				*data = NULL;
+			}
+			
+			if (this_types == NULL) {
+				for (i=1; i<self->nargs; i++) {
+					arg_types[i] = userdef;
+				}
+			}
+			else {
+				for (i=1; i<self->nargs; i++) {
+					arg_types[i] = this_types[i];
+				}
+			}
+			Py_DECREF(obj);
+			return 0;
 		}
-		if (obj == NULL) {
-			PyErr_SetString(PyExc_TypeError, 
-					"no registered loop for this "	\
-					"user-defined type");
-			return -1;			
-		}
-		if PyTuple_Check(obj) {
-			*function = (PyUFuncGenericFunction) \
-				PyCObject_AsVoidPtr(PyTuple_GET_ITEM(obj, 0));
-			*data = PyCObject_AsVoidPtr(PyTuple_GET_ITEM(obj, 1));
-		}
-		else {
-			*function = (PyUFuncGenericFunction)	\
-				PyCObject_AsVoidPtr(obj);
-			*data = NULL;
-		}
-                Py_DECREF(obj);
-		return 0;
 	}
-	
 
 	start_type = arg_types[0];
 	/* If the first argument is a scalar we need to place 
@@ -1210,13 +1224,15 @@ construct_matrices(PyUFuncLoopObject *loop, PyObject *args, PyArrayObject **mps)
 					scntcast += descr->elsize;
 				if (i < self->nin) {
 					loop->cast[i] = \
-						mps[i]->descr->f->cast[arg_types[i]];
+						PyArray_GetCastFunc(mps[i]->descr,
+								    arg_types[i]);
 				}
 				else {
-					loop->cast[i] = descr->f->	\
-						cast[mps[i]->descr->type_num];
+					loop->cast[i] =	PyArray_GetCastFunc \
+						(descr, mps[i]->descr->type_num);
 				}
 				Py_DECREF(descr);
+				if (!loop->cast[i]) return -1;
 			}
 			loop->swap[i] = !(PyArray_ISNOTSWAPPED(mps[i]));
 			if (loop->steps[i])
@@ -2993,6 +3009,7 @@ static int
 PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc, 
 			    int usertype,
 			    PyUFuncGenericFunction function,
+			    int *arg_types,
 			    void *data)
 {
 	PyArray_Descr *descr;
@@ -3002,7 +3019,7 @@ PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc,
 	descr=PyArray_DescrFromType(usertype);
 	if ((usertype < PyArray_USERDEF) || (descr==NULL)) {
 		PyErr_SetString(PyExc_TypeError, 
-				"unknown type");
+				"unknown user-defined type");
 		return -1;
 	}
 	Py_DECREF(descr);
@@ -3014,21 +3031,40 @@ PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc,
 	if (key == NULL) return -1;
 	cobj = PyCObject_FromVoidPtr((void *)function, NULL);
 	if (cobj == NULL) {Py_DECREF(key); return -1;}
-	if (data == NULL) {
+	if (data == NULL && arg_types == NULL) {
 		ret = PyDict_SetItem(ufunc->userloops, key, cobj);
 		Py_DECREF(cobj);
 		Py_DECREF(key);
 		return ret;
 	}
 	else {
-		PyObject *cobj2, *tmp;
-		cobj2 = PyCObject_FromVoidPtr(data, NULL);
-		if (cobj2 == NULL) {
-			Py_DECREF(cobj); 
-			Py_DECREF(key);
-			return -1;
+		PyObject *cobj2, *cobj3, *tmp;
+		if (arg_types == NULL) {
+			cobj2 = Py_None;
+			Py_INCREF(cobj2);
+		}
+		else {
+			cobj2 = PyCObject_FromVoidPtr((void *)arg_types, NULL);
+			if (cobj2 == NULL) {
+				Py_DECREF(cobj);
+				Py_DECREF(key);
+				return -1;
+			}
+		}
+		if (data == NULL) {
+			cobj3 = Py_None;
+			Py_INCREF(cobj3);
+		}
+		else {
+			cobj3 = PyCObject_FromVoidPtr(data, NULL);
+			if (cobj3 == NULL) {
+				Py_DECREF(cobj2);
+				Py_DECREF(cobj); 
+				Py_DECREF(key);
+				return -1;
+			}
 		}
-		tmp=Py_BuildValue("NN", cobj, cobj2);
+		tmp=Py_BuildValue("NNN", cobj, cobj2, cobj3);
 		ret = PyDict_SetItem(ufunc->userloops, key, tmp);
 		Py_DECREF(tmp);
 		Py_DECREF(key);