Merge pull request #2024 from lordmauve/better-swizzle

Fix OOB access in swizzle; add '0' and '1' swizzles

src_c/math.c

@@ -1584,6 +1584,7 @@ vector_str(pgVector *self)
 static PyObject *
 vector_getAttr_swizzle(pgVector *self, PyObject *attr_name)
 {
+    double value;
     double *coords;
     Py_ssize_t i, idx, len;
     PyObject *attr_unicode = NULL;
@@ -1634,23 +1635,33 @@ vector_getAttr_swizzle(pgVector *self, PyObject *attr_name)
             case 'y':
             case 'z':
                 idx = attr[i] - 'x';
-                break;
+                goto swizzle_idx;
             case 'w':
                 idx = 3;
+
+        swizzle_idx:
+                if (idx >= self->dim) {
+                    goto swizzle_failed;
+                };
+                value = coords[idx];
+                break;
+
+            case '0':
+                value = 0.0f;
                 break;
+            case '1':
+                value = 1.0f;
+                break;
+
             default:
                 goto swizzle_failed;
         }
         if (len == 2 || len == 3) {
-            ((pgVector *)res)->coords[i] = coords[idx];
-        }
-        else if (idx < self->dim) {
-            if (PyTuple_SetItem(res, i, PyFloat_FromDouble(coords[idx])) != 0)
+            ((pgVector *)res)->coords[i] = value;
+        } else {
+            if (PyTuple_SetItem(res, i, PyFloat_FromDouble(value)) != 0)
                 goto internal_error;
         }
-        else {
-            goto swizzle_failed;
-        }
     }
     /* swizzling succeeded! */
     Py_DECREF(attr_unicode);

test/math_test.py

@@ -2111,6 +2111,34 @@ class Vector(pygame.math.Vector3):
         v *= 2
         self.assertEqual(v, (4.0, 4.0, 4.0))
 
+    def test_swizzle_constants(self):
+        """We can get constant values from a swizzle."""
+        v = Vector2(7, 6)
+        self.assertEqual(
+            v.xy1,
+            (7.0, 6.0, 1.0),
+        )
+
+    def test_swizzle_four_constants(self):
+        """We can get 4 constant values from a swizzle."""
+        v = Vector2(7, 6)
+        self.assertEqual(
+            v.xy01,
+            (7.0, 6.0, 0.0, 1.0),
+        )
+
+    def test_swizzle_oob(self):
+        """An out-of-bounds swizzle raises an AttributeError."""
+        v = Vector2(7, 6)
+        with self.assertRaises(AttributeError):
+            v.xyz
+
+    def test_swizzle_set_oob(self):
+        """An out-of-bounds swizzle set raises an AttributeError."""
+        v = Vector2(7, 6)
+        with self.assertRaises(AttributeError):
+            v.xz = (1, 1)
+
 
 if __name__ == "__main__":
     unittest.main()