Allow configuring JPEG restart marker interval on save

libjpeg allows specifying the marker interval either in MCU blocks or in
MCU rows.  Support both, via separate parameters, rather than requiring
callers to do the math.

Co-authored-by: Andrew Murray <[email protected]>

Tests/test_file_jpeg.py

@@ -643,6 +643,26 @@ def test_save_low_quality_baseline_qtables(self):
             assert max(im2.quantization[0]) <= 255
             assert max(im2.quantization[1]) <= 255
 
+    @pytest.mark.parametrize(
+        "blocks, rows, markers",
+        ((0, 0, 0), (1, 0, 15), (3, 0, 5), (8, 0, 1), (0, 1, 3), (0, 2, 1)),
+    )
+    def test_restart_markers(self, blocks, rows, markers):
+        im = Image.new("RGB", (32, 32))  # 16 MCUs
+        out = BytesIO()
+        im.save(
+            out,
+            format="JPEG",
+            restart_marker_blocks=blocks,
+            restart_marker_rows=rows,
+            # force 8x8 pixel MCUs
+            subsampling=0,
+        )
+        markers_found = sum(
+            len(out.getvalue().split(bytes([0xFF, 0xD0 + i]))) - 1 for i in range(8)
+        )
+        assert markers_found == markers
+
     @pytest.mark.skipif(not djpeg_available(), reason="djpeg not available")
     def test_load_djpeg(self):
         with Image.open(TEST_FILE) as img:

docs/handbook/image-file-formats.rst

@@ -494,6 +494,18 @@ The :py:meth:`~PIL.Image.Image.save` method supports the following options:
 
     If absent, the setting will be determined by libjpeg or libjpeg-turbo.
 
+**restart_marker_blocks**
+    If present, emit a restart marker whenever the specified number of MCU
+    blocks has been produced.
+
+    .. versionadded:: 10.2.0
+
+**restart_marker_rows**
+    If present, emit a restart marker whenever the specified number of MCU
+    rows has been produced.
+
+    .. versionadded:: 10.2.0
+
 **qtables**
     If present, sets the qtables for the encoder. This is listed as an
     advanced option for wizards in the JPEG documentation. Use with

src/PIL/JpegImagePlugin.py

@@ -787,6 +787,8 @@ def validate_qtables(qtables):
         dpi[0],
         dpi[1],
         subsampling,
+        info.get("restart_marker_blocks", 0),
+        info.get("restart_marker_rows", 0),
         qtables,
         comment,
         extra,

src/encode.c

@@ -1045,6 +1045,8 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
     Py_ssize_t streamtype = 0; /* 0=interchange, 1=tables only, 2=image only */
     Py_ssize_t xdpi = 0, ydpi = 0;
     Py_ssize_t subsampling = -1; /* -1=default, 0=none, 1=medium, 2=high */
+    Py_ssize_t restart_marker_blocks = 0;
+    Py_ssize_t restart_marker_rows = 0;
     PyObject *qtables = NULL;
     unsigned int *qarrays = NULL;
     int qtablesLen = 0;
@@ -1057,7 +1059,7 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
 
     if (!PyArg_ParseTuple(
             args,
-            "ss|nnnnnnnnOz#y#y#",
+            "ss|nnnnnnnnnnOz#y#y#",
             &mode,
             &rawmode,
             &quality,
@@ -1068,6 +1070,8 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
             &xdpi,
             &ydpi,
             &subsampling,
+            &restart_marker_blocks,
+            &restart_marker_rows,
             &qtables,
             &comment,
             &comment_size,
@@ -1156,6 +1160,8 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
     ((JPEGENCODERSTATE *)encoder->state.context)->streamtype = streamtype;
     ((JPEGENCODERSTATE *)encoder->state.context)->xdpi = xdpi;
     ((JPEGENCODERSTATE *)encoder->state.context)->ydpi = ydpi;
+    ((JPEGENCODERSTATE *)encoder->state.context)->restart_marker_blocks = restart_marker_blocks;
+    ((JPEGENCODERSTATE *)encoder->state.context)->restart_marker_rows = restart_marker_rows;
     ((JPEGENCODERSTATE *)encoder->state.context)->comment = comment;
     ((JPEGENCODERSTATE *)encoder->state.context)->comment_size = comment_size;
     ((JPEGENCODERSTATE *)encoder->state.context)->extra = extra;

src/libImaging/Jpeg.h

@@ -83,6 +83,10 @@ typedef struct {
     /* Chroma Subsampling (-1=default, 0=none, 1=medium, 2=high) */
     int subsampling;
 
+    /* Restart marker interval, in MCU blocks or MCU rows, or 0 for none */
+    unsigned int restart_marker_blocks;
+    unsigned int restart_marker_rows;
+
     /* Converter input mode (input to the shuffler) */
     char rawmode[8 + 1];
 

src/libImaging/JpegEncode.c

@@ -210,6 +210,8 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             }
             context->cinfo.smoothing_factor = context->smooth;
             context->cinfo.optimize_coding = (boolean)context->optimize;
+            context->cinfo.restart_interval = context->restart_marker_blocks;
+            context->cinfo.restart_in_rows = context->restart_marker_rows;
             if (context->xdpi > 0 && context->ydpi > 0) {
                 context->cinfo.write_JFIF_header = TRUE;
                 context->cinfo.density_unit = 1; /* dots per inch */