Implement pitch argument for image.tobytes()

Author: whydoubt

buildconfig/stubs/pygame/image.pyi

@@ -20,6 +20,7 @@ def tostring(
     surface: Surface,
     format: _to_string_format,
     flipped: bool = False,
+    pitch: int = -1,
 ) -> bytes: ...
 def fromstring(
     bytes: bytes,
@@ -31,7 +32,10 @@ def fromstring(
 
 # the use of tobytes/frombytes is preferred over tostring/fromstring
 def tobytes(
-    surface: Surface, format: _to_string_format, flipped: bool = False
+    surface: Surface,
+    format: _to_string_format,
+    flipped: bool = False,
+    pitch: int = -1,
 ) -> bytes: ...
 def frombytes(
     bytes: bytes,

docs/reST/ref/image.rst

@@ -168,7 +168,7 @@ following formats.
 .. function:: tostring
 
    | :sl:`transfer image to byte buffer`
-   | :sg:`tostring(Surface, format, flipped=False) -> bytes`
+   | :sg:`tostring(Surface, format, flipped=False, pitch=-1) -> bytes`
 
    DEPRECATED: This function has the same functionality as :func:`tobytes()`, which is preferred and should be used.
 
@@ -179,7 +179,7 @@ following formats.
 .. function:: tobytes
 
    | :sl:`transfer image to byte buffer`
-   | :sg:`tobytes(Surface, format, flipped=False) -> bytes`
+   | :sg:`tobytes(Surface, format, flipped=False, pitch=-1) -> bytes`
 
    Creates a string of bytes that can be transferred with the ``fromstring``
    or ``frombytes`` methods in other Python imaging packages. Some Python
@@ -208,12 +208,19 @@ following formats.
 
       * ``ARGB_PREMULT``, 32-bit image with colors scaled by alpha channel, alpha channel first
 
+   The 'pitch' argument can be used to specify the pitch/stride per horizontal line
+   of the image in bytes. It must be equal to or greater than how many bytes
+   the pixel data of each horizontal line in the image bytes occupies without any
+   extra padding. By default, it is ``-1``, which means that the pitch/stride is
+   the same size as how many bytes the pure pixel data of each horizontal line takes.
+
    .. note:: The use of this function is recommended over :func:`tostring` as of pygame 2.1.3.
              This function was introduced so it matches nicely with other 
              libraries (PIL, numpy, etc), and with people's expectations.
 
    .. versionadded:: 2.1.3 
    .. versionchanged:: 2.2.0 Now supports keyword arguments.
+   .. versionchanged:: 2.4.0 Added a 'pitch' argument.
 
    .. ## pygame.image.tobytes ##
 

src_c/doc/image_doc.h

@@ -4,8 +4,8 @@
 #define DOC_IMAGE_SAVE "save(Surface, file) -> None\nsave(Surface, file, namehint="") -> None\nsave an image to file (or file-like object)"
 #define DOC_IMAGE_GETSDLIMAGEVERSION "get_sdl_image_version(linked=True) -> None\nget_sdl_image_version(linked=True) -> (major, minor, patch)\nget version number of the SDL_Image library being used"
 #define DOC_IMAGE_GETEXTENDED "get_extended() -> bool\ntest if extended image formats can be loaded"
-#define DOC_IMAGE_TOSTRING "tostring(Surface, format, flipped=False) -> bytes\ntransfer image to byte buffer"
-#define DOC_IMAGE_TOBYTES "tobytes(Surface, format, flipped=False) -> bytes\ntransfer image to byte buffer"
+#define DOC_IMAGE_TOSTRING "tostring(Surface, format, flipped=False, pitch=-1) -> bytes\ntransfer image to byte buffer"
+#define DOC_IMAGE_TOBYTES "tobytes(Surface, format, flipped=False, pitch=-1) -> bytes\ntransfer image to byte buffer"
 #define DOC_IMAGE_FROMSTRING "fromstring(bytes, size, format, flipped=False, pitch=-1) -> Surface\ncreate new Surface from a byte buffer"
 #define DOC_IMAGE_FROMBYTES "frombytes(bytes, size, format, flipped=False, pitch=-1) -> Surface\ncreate new Surface from a byte buffer"
 #define DOC_IMAGE_FROMBUFFER "frombuffer(buffer, size, format, pitch=-1) -> Surface\ncreate a new Surface that shares data inside a bytes buffer"

src_c/image.c

@@ -45,6 +45,12 @@ SaveTGA_RW(SDL_Surface *surface, SDL_RWops *out, int rle);
     ((flipped) ? (((char *)data) + (height - row - 1) * width) \
                : (((char *)data) + row * width))
 
+#define PITCHPAD(data, byte_width, pitch)    \
+    if (pitch > byte_width) {                \
+        memset(data, 0, pitch - byte_width); \
+        data += pitch - byte_width;          \
+    }
+
 static PyObject *extloadobj = NULL;
 static PyObject *extsaveobj = NULL;
 static PyObject *extverobj = NULL;
@@ -414,9 +420,10 @@ tobytes_surf_32bpp_sse42(SDL_Surface *surf, int flipped, char *data,
 static void
 tobytes_surf_32bpp(SDL_Surface *surf, int flipped, int hascolorkey,
                    Uint32 colorkey, char *serialized_image, int color_offset,
-                   int alpha_offset)
+                   int alpha_offset, int pitch)
 {
     int w, h;
+    int byte_width = surf->w * 4;
 
     Uint32 Rmask = surf->format->Rmask;
     Uint32 Gmask = surf->format->Gmask;
@@ -441,7 +448,8 @@ tobytes_surf_32bpp(SDL_Surface *surf, int flipped, int hascolorkey,
         sizeof(int) == sizeof(Uint32) &&
         4 * sizeof(Uint32) == sizeof(__m128i) &&
         !hascolorkey /* No color key */
-        && SDL_HasSSE42() == SDL_TRUE
+        && byte_width == pitch &&
+        SDL_HasSSE42() == SDL_TRUE
         /* The SSE code assumes it will always read at least 4 pixels */
         && surf->w >= 4
         /* Our SSE code assumes masks are at most 0xff */
@@ -480,6 +488,7 @@ tobytes_surf_32bpp(SDL_Surface *surf, int flipped, int hascolorkey,
                                    : 255);
             serialized_image += 4;
         }
+        PITCHPAD(serialized_image, byte_width, pitch)
     }
 }
 
@@ -490,25 +499,25 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
     PyObject *bytes = NULL;
     char *format, *data;
     SDL_Surface *surf;
-    int w, h, flipped = 0;
+    int w, h, flipped = 0, pitch = -1;
     int byte_width;
     Py_ssize_t len;
     Uint32 Rmask, Gmask, Bmask, Amask, Rshift, Gshift, Bshift, Ashift, Rloss,
         Gloss, Bloss, Aloss;
     int hascolorkey = 0;
     Uint32 color, colorkey;
     Uint32 alpha;
-    static char *kwds[] = {"surface", "format", "flipped", NULL};
+    static char *kwds[] = {"surface", "format", "flipped", "pitch", NULL};
 
 #ifdef _MSC_VER
     /* MSVC static analyzer false alarm: assure format is NULL-terminated by
      * making analyzer assume it was initialised */
     __analysis_assume(format = "inited");
 #endif
 
-    if (!PyArg_ParseTupleAndKeywords(arg, kwarg, "O!s|i", kwds,
+    if (!PyArg_ParseTupleAndKeywords(arg, kwarg, "O!s|ii", kwds,
                                      &pgSurface_Type, &surfobj, &format,
-                                     &flipped))
+                                     &flipped, &pitch))
         return NULL;
     surf = pgSurface_AsSurface(surfobj);
 
@@ -555,16 +564,28 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
         return RAISE(PyExc_ValueError, "Unrecognized type of format");
     }
 
-    bytes = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)byte_width * surf->h);
+    if (pitch == -1) {
+        pitch = byte_width;
+    }
+    else if (pitch < byte_width) {
+        return RAISE(PyExc_ValueError,
+                     "Pitch must be greater than or equal to the width "
+                     "as per the format");
+    }
+
+    bytes = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)pitch * surf->h);
     if (!bytes)
         return NULL;
     PyBytes_AsStringAndSize(bytes, &data, &len);
 
     if (!strcmp(format, "P")) {
         pgSurface_Lock(surfobj);
-        for (h = 0; h < surf->h; ++h)
-            memcpy(DATAROW(data, h, byte_width, surf->h, flipped),
-                   (char *)surf->pixels + (h * surf->pitch), surf->w);
+        for (h = 0; h < surf->h; ++h) {
+            Uint8 *ptr = (Uint8 *)DATAROW(data, h, pitch, surf->h, flipped);
+            memcpy(ptr, (char *)surf->pixels + (h * surf->pitch), surf->w);
+            if (pitch > byte_width)
+                memset(ptr + byte_width, 0, pitch - byte_width);
+        }
         pgSurface_Unlock(surfobj);
     }
     else if (!strcmp(format, "RGB")) {
@@ -582,6 +603,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[2] = (char)surf->format->palette->colors[color].b;
                         data += 3;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 2:
@@ -595,6 +617,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[2] = (char)(((color & Bmask) >> Bshift) << Bloss);
                         data += 3;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 3:
@@ -613,6 +636,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[2] = (char)(((color & Bmask) >> Bshift) << Bloss);
                         data += 3;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 4:
@@ -626,6 +650,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[2] = (char)(((color & Bmask) >> Bshift) << Rloss);
                         data += 3;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
         }
@@ -648,6 +673,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                                               : (char)255;
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 2:
@@ -667,6 +693,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                                                : 255);
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 3:
@@ -691,11 +718,12 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                                                : 255);
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 4:
                 tobytes_surf_32bpp(surf, flipped, hascolorkey, colorkey, data,
-                                   0, 3);
+                                   0, 3, pitch);
                 break;
         }
         pgSurface_Unlock(surfobj);
@@ -715,6 +743,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[0] = (char)255;
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 2:
@@ -731,6 +760,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                                                : 255);
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 3:
@@ -752,11 +782,12 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                                                : 255);
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 4:
                 tobytes_surf_32bpp(surf, flipped, hascolorkey, colorkey, data,
-                                   1, 0);
+                                   1, 0, pitch);
                 break;
         }
         pgSurface_Unlock(surfobj);
@@ -776,6 +807,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[3] = (char)255;
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 2:
@@ -792,6 +824,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                                                : 255);
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 3:
@@ -813,6 +846,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                                                : 255);
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 4:
@@ -829,6 +863,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                                                : 255);
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
         }
@@ -856,6 +891,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[3] = (char)alpha;
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 3:
@@ -882,6 +918,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[3] = (char)alpha;
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 4:
@@ -908,6 +945,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[3] = (char)alpha;
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
         }
@@ -935,6 +973,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[0] = (char)alpha;
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 3:
@@ -961,6 +1000,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[0] = (char)alpha;
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
             case 4:
@@ -987,6 +1027,7 @@ image_tobytes(PyObject *self, PyObject *arg, PyObject *kwarg)
                         data[0] = (char)alpha;
                         data += 4;
                     }
+                    PITCHPAD(data, byte_width, pitch)
                 }
                 break;
         }