Python: Add support for @ infix operator matrix multiplication

This differential revision implements the code for T56276

Reviewers: campbellbarton

Reviewed By: campbellbarton

Differential Revision: https://developer.blender.org/D3587

1 files changed, 129 insertions, 37 deletions

diff --git a/source/blender/python/mathutils/mathutils_Vector.c b/source/blender/python/mathutils/mathutils_Vector.c
index 6a40f22d9df..dc05f463d22 100644
--- a/source/blender/python/mathutils/mathutils_Vector.c
+++ b/source/blender/python/mathutils/mathutils_Vector.c
@@ -1706,12 +1706,25 @@ static PyObject *vector_mul_float(VectorObject *vec, const float scalar)
 	mul_vn_vn_fl(tvec, vec->vec, vec->size, scalar);
 	return Vector_CreatePyObject_alloc(tvec, vec->size, Py_TYPE(vec));
 }
+#ifdef USE_MATHUTILS_ELEM_MUL
+static PyObject *vector_mul_vec(VectorObject *vec1, VectorObject *vec2)
+{
+	float *tvec = PyMem_Malloc(vec1->size * sizeof(float));
+	if (tvec == NULL) {
+		PyErr_SetString(PyExc_MemoryError,
+						"vec * vec: "
+						"problem allocating pointer space");
+		return NULL;
+	}
 
+	mul_vn_vnvn(tvec, vec1->vec, vec2->vec, vec1->size);
+	return Vector_CreatePyObject_alloc(tvec, vec1->size, Py_TYPE(vec1));
+}
+#endif
 static PyObject *Vector_mul(PyObject *v1, PyObject *v2)
 {
 	VectorObject *vec1 = NULL, *vec2 = NULL;
 	float scalar;
-	int vec_size;
 
 	if (VectorObject_Check(v1)) {
 		vec1 = (VectorObject *)v1;
@@ -1729,6 +1742,7 @@ static PyObject *Vector_mul(PyObject *v1, PyObject *v2)
 
 	/* make sure v1 is always the vector */
 	if (vec1 && vec2) {
+#ifdef USE_MATHUTILS_ELEM_MUL
 		if (vec1->size != vec2->size) {
 			PyErr_SetString(PyExc_ValueError,
 			                "Vector multiplication: "
@@ -1736,30 +1750,12 @@ static PyObject *Vector_mul(PyObject *v1, PyObject *v2)
 			return NULL;
 		}
 
-		/*dot product*/
-		return PyFloat_FromDouble(dot_vn_vn(vec1->vec, vec2->vec, vec1->size));
+		/* element-wise product */
+		return vector_mul_vec(vec1, vec2);
+#endif
 	}
 	else if (vec1) {
-		if (MatrixObject_Check(v2)) {
-			/* VEC * MATRIX */
-			float tvec[MAX_DIMENSIONS];
-
-			if (BaseMath_ReadCallback((MatrixObject *)v2) == -1)
-				return NULL;
-			if (row_vector_multiplication(tvec, vec1, (MatrixObject *)v2) == -1) {
-				return NULL;
-			}
-
-			if (((MatrixObject *)v2)->num_row == 4 && vec1->size == 3) {
-				vec_size = 3;
-			}
-			else {
-				vec_size = ((MatrixObject *)v2)->num_col;
-			}
-
-			return Vector_CreatePyObject(tvec, vec_size, Py_TYPE(vec1));
-		}
-		else if (((scalar = PyFloat_AsDouble(v2)) == -1.0f && PyErr_Occurred()) == 0) { /* VEC * FLOAT */
+		if (((scalar = PyFloat_AsDouble(v2)) == -1.0f && PyErr_Occurred()) == 0) { /* VEC * FLOAT */
 			return vector_mul_float(vec1, scalar);
 		}
 	}
@@ -1768,12 +1764,9 @@ static PyObject *Vector_mul(PyObject *v1, PyObject *v2)
 			return vector_mul_float(vec2, scalar);
 		}
 	}
-	else {
-		BLI_assert(!"internal error");
-	}
 
 	PyErr_Format(PyExc_TypeError,
-	             "Vector multiplication: "
+				 "Element-wise multiplication: "
 	             "not supported between '%.200s' and '%.200s' types",
 	             Py_TYPE(v1)->tp_name, Py_TYPE(v2)->tp_name);
 	return NULL;
@@ -1782,32 +1775,129 @@ static PyObject *Vector_mul(PyObject *v1, PyObject *v2)
 /* multiplication in-place: obj *= obj */
 static PyObject *Vector_imul(PyObject *v1, PyObject *v2)
 {
-	VectorObject *vec = (VectorObject *)v1;
+	VectorObject *vec1 = NULL, *vec2 = NULL;
 	float scalar;
 
-	if (BaseMath_ReadCallback_ForWrite(vec) == -1)
+	if (VectorObject_Check(v1)) {
+		vec1 = (VectorObject *)v1;
+		if (BaseMath_ReadCallback(vec1) == -1)
+			return NULL;
+	}
+	if (VectorObject_Check(v2)) {
+		vec2 = (VectorObject *)v2;
+		if (BaseMath_ReadCallback(vec2) == -1)
+			return NULL;
+	}
+
+	if (BaseMath_ReadCallback_ForWrite(vec1) == -1)
 		return NULL;
 
 	/* Intentionally don't support (Quaternion, Matrix) here, uses reverse order instead. */
 
-	/* only support 'vec *= float'
-	 *  vec*=vec result is a float so that wont work */
-	if (((scalar = PyFloat_AsDouble(v2)) == -1.0f && PyErr_Occurred()) == 0) { /* VEC *= FLOAT */
-		mul_vn_fl(vec->vec, vec->size, scalar);
+	if (vec1 && vec2) {
+#ifdef USE_MATHUTILS_ELEM_MUL
+		if (vec1->size != vec2->size) {
+			PyErr_SetString(PyExc_ValueError,
+							"Vector multiplication: "
+							"vectors must have the same dimensions for this operation");
+			return NULL;
+		}
+
+		/* element-wise product inplace */
+		mul_vn_vn(vec1->vec, vec2->vec, vec1->size);
+#else
+		PyErr_Format(PyExc_TypeError,
+					 "Inplace element-wise multiplication: "
+					 "not supported between '%.200s' and '%.200s' types",
+					 Py_TYPE(v1)->tp_name, Py_TYPE(v2)->tp_name);
+		return NULL;
+#endif
+	}
+	else if (vec1 && (((scalar = PyFloat_AsDouble(v2)) == -1.0f && PyErr_Occurred()) == 0)) { /* VEC *= FLOAT */
+		mul_vn_fl(vec1->vec, vec1->size, scalar);
 	}
 	else {
 		PyErr_Format(PyExc_TypeError,
-		             "Vector multiplication: (%s *= %s) "
-		             "invalid type for this operation",
-		             Py_TYPE(v1)->tp_name, Py_TYPE(v2)->tp_name);
+					 "Inplace element-wise multiplication: "
+					 "not supported between '%.200s' and '%.200s' types",
+					 Py_TYPE(v1)->tp_name, Py_TYPE(v2)->tp_name);
 		return NULL;
 	}
 
-	(void)BaseMath_WriteCallback(vec);
+	(void)BaseMath_WriteCallback(vec1);
 	Py_INCREF(v1);
 	return v1;
 }
 
+static PyObject *Vector_matmul(PyObject *v1, PyObject *v2)
+{
+	VectorObject *vec1 = NULL, *vec2 = NULL;
+	int vec_size;
+
+	if (VectorObject_Check(v1)) {
+		vec1 = (VectorObject *)v1;
+		if (BaseMath_ReadCallback(vec1) == -1)
+			return NULL;
+	}
+	if (VectorObject_Check(v2)) {
+		vec2 = (VectorObject *)v2;
+		if (BaseMath_ReadCallback(vec2) == -1)
+			return NULL;
+	}
+
+
+	/* Intentionally don't support (Quaternion) here, uses reverse order instead. */
+
+	/* make sure v1 is always the vector */
+	if (vec1 && vec2) {
+		if (vec1->size != vec2->size) {
+			PyErr_SetString(PyExc_ValueError,
+							"Vector multiplication: "
+							"vectors must have the same dimensions for this operation");
+			return NULL;
+		}
+
+		/*dot product*/
+		return PyFloat_FromDouble(dot_vn_vn(vec1->vec, vec2->vec, vec1->size));
+	}
+	else if (vec1) {
+		if (MatrixObject_Check(v2)) {
+			/* VEC @ MATRIX */
+			float tvec[MAX_DIMENSIONS];
+
+			if (BaseMath_ReadCallback((MatrixObject *)v2) == -1)
+				return NULL;
+			if (row_vector_multiplication(tvec, vec1, (MatrixObject *)v2) == -1) {
+				return NULL;
+			}
+
+			if (((MatrixObject *)v2)->num_row == 4 && vec1->size == 3) {
+				vec_size = 3;
+			}
+			else {
+				vec_size = ((MatrixObject *)v2)->num_col;
+			}
+
+			return Vector_CreatePyObject(tvec, vec_size, Py_TYPE(vec1));
+		}
+	}
+
+	PyErr_Format(PyExc_TypeError,
+				 "Vector multiplication: "
+				 "not supported between '%.200s' and '%.200s' types",
+				 Py_TYPE(v1)->tp_name, Py_TYPE(v2)->tp_name);
+	return NULL;
+}
+
+static PyObject *Vector_imatmul(PyObject *v1, PyObject *v2)
+{
+	PyErr_Format(PyExc_TypeError,
+				 "Inplace vector multiplication: "
+				 "not supported between '%.200s' and '%.200s' types",
+				 Py_TYPE(v1)->tp_name, Py_TYPE(v2)->tp_name);
+	return NULL;
+}
+
 /* divid: obj / obj */
 static PyObject *Vector_div(PyObject *v1, PyObject *v2)
 {
@@ -2119,6 +2209,8 @@ static PyNumberMethods Vector_NumMethods = {
 	NULL,                       /* nb_inplace_floor_divide */
 	Vector_idiv,                /* nb_inplace_true_divide */
 	NULL,                       /* nb_index */
+	(binaryfunc) Vector_matmul, /* nb_matrix_multiply */
+	(binaryfunc) Vector_imatmul, /* nb_inplace_matrix_multiply */
 };
 
 /*------------------PY_OBECT DEFINITION--------------------------*/