From 4d7372bfd0293fc133f7bf29ce17d2e7b75a8c9a Mon Sep 17 00:00:00 2001
From: Benjamin Gilbert <bgilbert@backtick.net>
Date: Tue, 24 Oct 2023 00:05:30 -0500
Subject: [PATCH] Implement streamtype=1 option for tables-only JPEG encoding

We already support streamtype=2 to skip producing JPEG tables, but
streamtype=1, which skips everything but the tables, was never implemented.
The streamtype=1 stub code dates to Git pre-history, so it's not
immediately clear why.  Implement the missing support.

jpeg_write_tables() can't resume after a full output buffer (it fails with
JERR_CANT_SUSPEND), so it might seem that Pillow needs to pre-compute the
necessary buffer size.  However, in the normal case of producing an
interchange stream, the tables are written via the same libjpeg codepath
during the first jpeg_write_scanlines() call, and table writes aren't
resumable there either.  Thus, any buffer large enough for the normal case
will also be large enough for a tables-only file.

The streamtype option isn't documented and this commit doesn't change that.
It does add a test though.

Co-authored-by: Andrew Murray <radarhere@users.noreply.github.com>
---
 Tests/test_file_jpeg.py     | 22 ++++++++++++++++++++++
 src/libImaging/JpegEncode.c |  7 ++++---
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/Tests/test_file_jpeg.py b/Tests/test_file_jpeg.py
index a0822d000..b87cbf0c7 100644
--- a/Tests/test_file_jpeg.py
+++ b/Tests/test_file_jpeg.py
@@ -961,6 +961,28 @@ class TestFileJpeg:
             im.load()
             ImageFile.LOAD_TRUNCATED_IMAGES = False
 
+    def test_separate_tables(self):
+        im = hopper()
+        data = []  # [interchange, tables-only, image-only]
+        for streamtype in range(3):
+            out = BytesIO()
+            im.save(out, format="JPEG", streamtype=streamtype)
+            data.append(out.getvalue())
+
+        # SOI, EOI
+        for marker in b"\xff\xd8", b"\xff\xd9":
+            assert marker in data[1] and marker in data[2]
+        # DHT, DQT
+        for marker in b"\xff\xc4", b"\xff\xdb":
+            assert marker in data[1] and marker not in data[2]
+        # SOF0, SOS, APP0 (JFIF header)
+        for marker in b"\xff\xc0", b"\xff\xda", b"\xff\xe0":
+            assert marker not in data[1] and marker in data[2]
+
+        with Image.open(BytesIO(data[0])) as interchange_im:
+            with Image.open(BytesIO(data[1] + data[2])) as combined_im:
+                assert_image_equal(interchange_im, combined_im)
+
     def test_repr_jpeg(self):
         im = hopper()
 
diff --git a/src/libImaging/JpegEncode.c b/src/libImaging/JpegEncode.c
index 2a24eff39..5ebcea6f6 100644
--- a/src/libImaging/JpegEncode.c
+++ b/src/libImaging/JpegEncode.c
@@ -218,9 +218,9 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             }
             switch (context->streamtype) {
                 case 1:
-                    /* tables only -- not yet implemented */
-                    state->errcode = IMAGING_CODEC_CONFIG;
-                    return -1;
+                    /* tables only */
+                    jpeg_write_tables(&context->cinfo);
+                    goto cleanup;
                 case 2:
                     /* image only */
                     jpeg_suppress_tables(&context->cinfo, TRUE);
@@ -316,6 +316,7 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             }
             jpeg_finish_compress(&context->cinfo);
 
+cleanup:
             /* Clean up */
             if (context->comment) {
                 free(context->comment);