Merge pull request #2111 from ScriptLineStudios/Surface.set_at

vector2 support for Surface.get_at and set_at

Author: Starbuck5

buildconfig/stubs/pygame/surface.pyi

@@ -7,7 +7,6 @@ from pygame.rect import FRect, Rect
 from ._common import (
     ColorValue,
     Coordinate,
-    IntCoordinate,
     Literal,
     RectValue,
     RGBAOutput,
@@ -117,9 +116,9 @@ class Surface:
     def mustlock(self) -> bool: ...
     def get_locked(self) -> bool: ...
     def get_locks(self) -> Tuple[Any, ...]: ...
-    def get_at(self, x_y: IntCoordinate) -> Color: ...
-    def set_at(self, x_y: IntCoordinate, color: ColorValue) -> None: ...
-    def get_at_mapped(self, x_y: IntCoordinate) -> int: ...
+    def get_at(self, x_y: Coordinate) -> Color: ...
+    def set_at(self, x_y: Coordinate, color: ColorValue) -> None: ...
+    def get_at_mapped(self, x_y: Coordinate) -> int: ...
     def get_palette(self) -> List[Color]: ...
     def get_palette_at(self, index: int) -> Color: ...
     def set_palette(self, palette: Sequence[ColorValue]) -> None: ...

docs/reST/ref/surface.rst

@@ -445,12 +445,15 @@
       methods - or by using :mod:`pygame.surfarray`/:mod:`pygame.PixelArray`.
 
       This function will temporarily lock and unlock the Surface as needed.
-
-      .. versionaddedold:: 1.9
+      
+      .. versionchanged:: 2.3.1 can now also accept both float coordinates and Vector2s for pixels.
+      
          Returning a Color instead of tuple. Use ``tuple(surf.get_at((x,y)))``
          if you want a tuple, and not a Color. This should only matter if
          you want to use the color as a key in a dict.
-
+      
+      .. versionaddedold:: 1.9
+      
       .. ## Surface.get_at ##
 
    .. method:: set_at
@@ -471,6 +474,8 @@
       .. note:: If the surface is palettized, the pixel color will be set to the
                 most similar color in the palette.
 
+      .. versionchanged:: 2.3.1 can now also accept both float coordinates and Vector2s for pixels.
+
       .. ## Surface.set_at ##
 
    .. method:: get_at_mapped
@@ -489,6 +494,8 @@
 
       .. versionaddedold:: 1.9.2
 
+      .. versionchanged:: 2.3.1 can now also accept both float coordinates and Vector2s for pixels.
+
       .. ## Surface.get_at_mapped ##
 
    .. method:: get_palette

src_c/surface.c

@@ -304,10 +304,9 @@ static PyGetSetDef surface_getsets[] = {
     {NULL, NULL, NULL, NULL, NULL}};
 
 static struct PyMethodDef surface_methods[] = {
-    {"get_at", surf_get_at, METH_VARARGS, DOC_SURFACE_GETAT},
+    {"get_at", surf_get_at, METH_O, DOC_SURFACE_GETAT},
     {"set_at", surf_set_at, METH_VARARGS, DOC_SURFACE_SETAT},
-    {"get_at_mapped", surf_get_at_mapped, METH_VARARGS,
-     DOC_SURFACE_GETATMAPPED},
+    {"get_at_mapped", surf_get_at_mapped, METH_O, DOC_SURFACE_GETATMAPPED},
     {"map_rgb", surf_map_rgb, METH_VARARGS, DOC_SURFACE_MAPRGB},
     {"unmap_rgb", surf_unmap_rgb, METH_O, DOC_SURFACE_UNMAPRGB},
 
@@ -774,7 +773,7 @@ _raise_create_surface_error(void)
 
 /* surface object methods */
 static PyObject *
-surf_get_at(PyObject *self, PyObject *args)
+surf_get_at(PyObject *self, PyObject *position)
 {
     SDL_Surface *surf = pgSurface_AsSurface(self);
     SDL_PixelFormat *format = NULL;
@@ -784,10 +783,13 @@ surf_get_at(PyObject *self, PyObject *args)
     Uint8 *pix;
     Uint8 rgba[4] = {0, 0, 0, 255};
 
-    if (!PyArg_ParseTuple(args, "(ii)", &x, &y))
-        return NULL;
     SURF_INIT_CHECK(surf)
 
+    if (!pg_TwoIntsFromObj(position, &x, &y)) {
+        return RAISE(PyExc_TypeError,
+                     "position must be a sequence of two numbers");
+    }
+
     if (x < 0 || x >= surf->w || y < 0 || y >= surf->h)
         return RAISE(PyExc_IndexError, "pixel index out of range");
 
@@ -838,17 +840,23 @@ surf_set_at(PyObject *self, PyObject *args)
     SDL_PixelFormat *format = NULL;
     Uint8 *pixels;
     int x, y;
+    PyObject *position;
     Uint32 color;
     Uint8 rgba[4] = {0, 0, 0, 0};
     PyObject *rgba_obj;
     Uint8 *byte_buf;
 
-    if (!PyArg_ParseTuple(args, "(ii)O", &x, &y, &rgba_obj))
+    if (!PyArg_ParseTuple(args, "OO", &position, &rgba_obj))
         return NULL;
     SURF_INIT_CHECK(surf)
 
     format = surf->format;
 
+    if (!pg_TwoIntsFromObj(position, &x, &y)) {
+        return RAISE(PyExc_TypeError,
+                     "position must be a sequence of two numbers");
+    }
+
     if (format->BytesPerPixel < 1 || format->BytesPerPixel > 4)
         return RAISE(PyExc_RuntimeError, "invalid color depth for surface");
 
@@ -913,7 +921,7 @@ surf_set_at(PyObject *self, PyObject *args)
 }
 
 static PyObject *
-surf_get_at_mapped(PyObject *self, PyObject *args)
+surf_get_at_mapped(PyObject *self, PyObject *position)
 {
     SDL_Surface *surf = pgSurface_AsSurface(self);
     SDL_PixelFormat *format = NULL;
@@ -922,10 +930,13 @@ surf_get_at_mapped(PyObject *self, PyObject *args)
     Sint32 color;
     Uint8 *pix;
 
-    if (!PyArg_ParseTuple(args, "(ii)", &x, &y))
-        return NULL;
     SURF_INIT_CHECK(surf)
 
+    if (!pg_TwoIntsFromObj(position, &x, &y)) {
+        return RAISE(PyExc_TypeError,
+                     "position must be a sequence of two numbers");
+    }
+
     if (x < 0 || x >= surf->w || y < 0 || y >= surf->h)
         return RAISE(PyExc_IndexError, "pixel index out of range");
 

test/surface_test.py

@@ -114,6 +114,11 @@ def test_set_at(self):
         self.assertIsInstance(r, pygame.Color)
         self.assertEqual(r, (10, 10, 10, 255))
 
+        s.set_at(pygame.Vector2(0, 0), (10, 10, 10, 255))
+        r = s.get_at(pygame.Vector2(0, 0))
+        self.assertIsInstance(r, pygame.Color)
+        self.assertEqual(r, (10, 10, 10, 255))
+
         # try setting a color with a single integer.
         s.fill((0, 0, 0, 255))
         s.set_at((10, 1), 0x0000FF)
@@ -1905,10 +1910,15 @@ def test_get_at(self):
         surf.set_at((1, 1), c11)
         c = surf.get_at((0, 0))
         self.assertIsInstance(c, pygame.Color)
+        c = surf.get_at(pygame.Vector2(0, 0))
+        self.assertIsInstance(c, pygame.Color)
         self.assertEqual(c, c00)
         self.assertEqual(surf.get_at((0, 1)), c01)
         self.assertEqual(surf.get_at((1, 0)), c10)
         self.assertEqual(surf.get_at((1, 1)), c11)
+        self.assertEqual(surf.get_at(pygame.Vector2(0, 1)), c01)
+        self.assertEqual(surf.get_at(pygame.Vector2(1, 0)), c10)
+        self.assertEqual(surf.get_at(pygame.Vector2(1, 1)), c11)
         for p in [(-1, 0), (0, -1), (2, 0), (0, 2)]:
             self.assertRaises(IndexError, surf.get_at, p)
 
@@ -1923,6 +1933,12 @@ def test_get_at_mapped(self):
                 surf.map_rgb(color),
                 "%i != %i, bitsize: %i" % (pixel, surf.map_rgb(color), bitsize),
             )
+            pixel = surf.get_at_mapped(pygame.Vector2(0, 0))
+            self.assertEqual(
+                pixel,
+                surf.map_rgb(color),
+                "%i != %i, bitsize: %i" % (pixel, surf.map_rgb(color), bitsize),
+            )
 
     def test_get_bitsize(self):
         pygame.display.init()