Merge pull request #3408 from Matiiss/matiiss-fix-vector-init-from-numpy-arrays

Fix `pygame.Vector{2,3}` initialization from numpy arrays

Author: damusss

src_c/math.c

@@ -2280,19 +2280,23 @@ static int
 _vector2_set(pgVector *self, PyObject *xOrSequence, PyObject *y)
 {
     if (xOrSequence) {
-        if (RealNumber_Check(xOrSequence)) {
-            self->coords[0] = PyFloat_AsDouble(xOrSequence);
-            /* scalar constructor. */
-            if (y == NULL) {
-                self->coords[1] = self->coords[0];
+        if (pgVectorCompatible_Check(xOrSequence, self->dim)) {
+            if (!PySequence_AsVectorCoords(xOrSequence, self->coords, 2)) {
+                return -1;
+            }
+            else {
                 return 0;
             }
         }
-        else if (pgVectorCompatible_Check(xOrSequence, self->dim)) {
-            if (!PySequence_AsVectorCoords(xOrSequence, self->coords, 2)) {
+        else if (RealNumber_Check(xOrSequence)) {
+            self->coords[0] = PyFloat_AsDouble(xOrSequence);
+            if (self->coords[0] == -1.0 && PyErr_Occurred()) {
                 return -1;
             }
-            else {
+
+            /* scalar constructor. */
+            if (y == NULL) {
+                self->coords[1] = self->coords[0];
                 return 0;
             }
         }
@@ -2323,6 +2327,9 @@ _vector2_set(pgVector *self, PyObject *xOrSequence, PyObject *y)
 
     if (RealNumber_Check(y)) {
         self->coords[1] = PyFloat_AsDouble(y);
+        if (self->coords[1] == -1.0 && PyErr_Occurred()) {
+            return -1;
+        }
     }
     else {
         goto error;
@@ -2718,20 +2725,24 @@ static int
 _vector3_set(pgVector *self, PyObject *xOrSequence, PyObject *y, PyObject *z)
 {
     if (xOrSequence) {
-        if (RealNumber_Check(xOrSequence)) {
-            self->coords[0] = PyFloat_AsDouble(xOrSequence);
-            /* scalar constructor. */
-            if (y == NULL && z == NULL) {
-                self->coords[1] = self->coords[0];
-                self->coords[2] = self->coords[0];
+        if (pgVectorCompatible_Check(xOrSequence, self->dim)) {
+            if (!PySequence_AsVectorCoords(xOrSequence, self->coords, 3)) {
+                return -1;
+            }
+            else {
                 return 0;
             }
         }
-        else if (pgVectorCompatible_Check(xOrSequence, self->dim)) {
-            if (!PySequence_AsVectorCoords(xOrSequence, self->coords, 3)) {
+        else if (RealNumber_Check(xOrSequence)) {
+            self->coords[0] = PyFloat_AsDouble(xOrSequence);
+            if (self->coords[0] == -1.0 && PyErr_Occurred()) {
                 return -1;
             }
-            else {
+
+            /* scalar constructor. */
+            if (y == NULL && z == NULL) {
+                self->coords[1] = self->coords[0];
+                self->coords[2] = self->coords[0];
                 return 0;
             }
         }
@@ -2764,7 +2775,14 @@ _vector3_set(pgVector *self, PyObject *xOrSequence, PyObject *y, PyObject *z)
     else if (y && z) {
         if (RealNumber_Check(y) && RealNumber_Check(z)) {
             self->coords[1] = PyFloat_AsDouble(y);
+            if (self->coords[1] == -1.0 && PyErr_Occurred()) {
+                return -1;
+            }
+
             self->coords[2] = PyFloat_AsDouble(z);
+            if (self->coords[2] == -1.0 && PyErr_Occurred()) {
+                return -1;
+            }
         }
         else {
             goto error;

test/math_test.py

@@ -6,6 +6,11 @@
 import pygame.math
 from pygame.math import Vector2, Vector3
 
+try:
+    import numpy
+except ModuleNotFoundError:
+    numpy = None
+
 IS_PYPY = "PyPy" == platform.python_implementation()
 
 
@@ -255,6 +260,48 @@ def testConstructionVector2(self):
         self.assertEqual(v.x, 1.2)
         self.assertEqual(v.y, 3.4)
 
+    def testConstructionNumericSequence(self):
+        class NumericSequence:
+            # PyFloat_AsDouble will use this to convert to a float
+            # so this is testing the implementation a bit
+            def __float__(self):
+                raise TypeError("Cannot convert to float")
+
+            def __getitem__(self, index):
+                return [1, 0][index]
+
+            def __len__(self):
+                return 2
+
+        v = Vector2(NumericSequence())
+        self.assertEqual(v.x, 1.0)
+        self.assertEqual(v.y, 0.0)
+
+    def testConstructionNumericNonFloat(self):
+        class NumericNonFloat:
+            # PyFloat_AsDouble will use this to convert to a float
+            # so this is testing the implementation a bit
+            def __float__(self):
+                raise TypeError("Cannot convert to float")
+
+        with self.assertRaises(TypeError):
+            Vector2(NumericNonFloat())
+
+        with self.assertRaises(TypeError):
+            Vector2(NumericNonFloat(), NumericNonFloat())
+
+        with self.assertRaises(TypeError):
+            Vector2(1.0, NumericNonFloat())
+
+    @unittest.skipIf(numpy is None, "numpy not available")
+    def testConstructionNumpyArray(self):
+        assert numpy is not None
+
+        arr = numpy.array([1.2, 3.4])
+        v = Vector2(arr)
+        self.assertEqual(v.x, 1.2)
+        self.assertEqual(v.y, 3.4)
+
     def testAttributeAccess(self):
         tmp = self.v1.x
         self.assertEqual(tmp, self.v1.x)
@@ -1431,6 +1478,53 @@ def testConstructionMissing(self):
         self.assertRaises(ValueError, Vector3, 1, 2)
         self.assertRaises(ValueError, Vector3, x=1, y=2)
 
+    def testConstructionNumericSequence(self):
+        class NumericSequence:
+            # PyFloat_AsDouble will use this to convert to a float
+            # so this is testing the implementation a bit
+            def __float__(self):
+                raise TypeError("Cannot convert to float")
+
+            def __getitem__(self, index):
+                return [1, 0, 5][index]
+
+            def __len__(self):
+                return 3
+
+        v = Vector3(NumericSequence())
+        self.assertEqual(v.x, 1.0)
+        self.assertEqual(v.y, 0.0)
+        self.assertEqual(v.z, 5.0)
+
+    def testConstructionNumericNonFloat(self):
+        class NumericNonFloat:
+            # PyFloat_AsDouble will use this to convert to a float
+            # so this is testing the implementation a bit
+            def __float__(self):
+                raise TypeError("Cannot convert to float")
+
+        with self.assertRaises(TypeError):
+            Vector3(NumericNonFloat())
+
+        with self.assertRaises(TypeError):
+            Vector3(NumericNonFloat(), NumericNonFloat(), NumericNonFloat())
+
+        with self.assertRaises(TypeError):
+            Vector3(1.0, NumericNonFloat(), 5.0)
+
+        with self.assertRaises(TypeError):
+            Vector3(1.0, 0.0, NumericNonFloat())
+
+    @unittest.skipIf(numpy is None, "numpy not available")
+    def testConstructionNumpyArray(self):
+        assert numpy is not None
+
+        arr = numpy.array([1.2, 3.4, 5.6], dtype=float)
+        v = Vector3(arr)
+        self.assertEqual(v.x, 1.2)
+        self.assertEqual(v.y, 3.4)
+        self.assertEqual(v.z, 5.6)
+
     def testAttributeAccess(self):
         tmp = self.v1.x
         self.assertEqual(tmp, self.v1.x)