Merge pull request #6774 from smason/write-jpeg-com

Support saving JPEG comments

Tests/test_file_jpeg.py

@@ -86,6 +86,33 @@ class TestFileJpeg:
             assert len(im.applist) == 2
 
             assert im.info["comment"] == b"File written by Adobe Photoshop\xa8 4.0\x00"
+            assert im.app["COM"] == im.info["comment"]
+
+    def test_comment_write(self):
+        with Image.open(TEST_FILE) as im:
+            assert im.info["comment"] == b"File written by Adobe Photoshop\xa8 4.0\x00"
+
+            # Test that existing comment is saved by default
+            out = BytesIO()
+            im.save(out, format="JPEG")
+            with Image.open(out) as reloaded:
+                assert im.info["comment"] == reloaded.info["comment"]
+
+            # Ensure that a blank comment causes any existing comment to be removed
+            for comment in ("", b"", None):
+                out = BytesIO()
+                im.save(out, format="JPEG", comment=comment)
+                with Image.open(out) as reloaded:
+                    assert "comment" not in reloaded.info
+
+            # Test that a comment argument overrides the default comment
+            for comment in ("Test comment text", b"Text comment text"):
+                out = BytesIO()
+                im.save(out, format="JPEG", comment=comment)
+                with Image.open(out) as reloaded:
+                    if not isinstance(comment, bytes):
+                        comment = comment.encode()
+                    assert reloaded.info["comment"] == comment
 
     def test_cmyk(self):
         # Test CMYK handling.  Thanks to Tim and Charlie for test data,

src/PIL/JpegImagePlugin.py

@@ -45,6 +45,7 @@ from . import Image, ImageFile, TiffImagePlugin
 from ._binary import i16be as i16
 from ._binary import i32be as i32
 from ._binary import o8
+from ._binary import o16be as o16
 from ._deprecate import deprecate
 from .JpegPresets import presets
 
@@ -724,7 +725,7 @@ def _save(im, fp, filename):
             icc_profile = icc_profile[MAX_DATA_BYTES_IN_MARKER:]
         i = 1
         for marker in markers:
-            size = struct.pack(">H", 2 + ICC_OVERHEAD_LEN + len(marker))
+            size = o16(2 + ICC_OVERHEAD_LEN + len(marker))
             extra += (
                 b"\xFF\xE2"
                 + size
@@ -735,6 +736,8 @@ def _save(im, fp, filename):
             )
             i += 1
 
+    comment = info.get("comment", im.info.get("comment"))
+
     # "progressive" is the official name, but older documentation
     # says "progression"
     # FIXME: issue a warning if the wrong form is used (post-1.1.7)
@@ -757,6 +760,7 @@ def _save(im, fp, filename):
         dpi[1],
         subsampling,
         qtables,
+        comment,
         extra,
         exif,
     )

src/encode.c

@@ -1048,6 +1048,8 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
     PyObject *qtables = NULL;
     unsigned int *qarrays = NULL;
     int qtablesLen = 0;
+    char *comment = NULL;
+    Py_ssize_t comment_size;
     char *extra = NULL;
     Py_ssize_t extra_size;
     char *rawExif = NULL;
@@ -1055,7 +1057,7 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
 
     if (!PyArg_ParseTuple(
             args,
-            "ss|nnnnnnnnOy#y#",
+            "ss|nnnnnnnnOz#y#y#",
             &mode,
             &rawmode,
             &quality,
@@ -1067,6 +1069,8 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
             &ydpi,
             &subsampling,
             &qtables,
+            &comment,
+            &comment_size,
             &extra,
             &extra_size,
             &rawExif,
@@ -1090,13 +1094,28 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
         return NULL;
     }
 
-    // Freed in JpegEncode, Case 5
+    // Freed in JpegEncode, Case 6
     qarrays = get_qtables_arrays(qtables, &qtablesLen);
 
+    if (comment && comment_size > 0) {
+        /* malloc check ok, length is from python parsearg */
+        char *p = malloc(comment_size);  // Freed in JpegEncode, Case 6
+        if (!p) {
+            return ImagingError_MemoryError();
+        }
+        memcpy(p, comment, comment_size);
+        comment = p;
+    } else {
+        comment = NULL;
+    }
+
     if (extra && extra_size > 0) {
         /* malloc check ok, length is from python parsearg */
-        char *p = malloc(extra_size);  // Freed in JpegEncode, Case 5
+        char *p = malloc(extra_size);  // Freed in JpegEncode, Case 6
         if (!p) {
+            if (comment) {
+                free(comment);
+            }
             return ImagingError_MemoryError();
         }
         memcpy(p, extra, extra_size);
@@ -1107,8 +1126,11 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
 
     if (rawExif && rawExifLen > 0) {
         /* malloc check ok, length is from python parsearg */
-        char *pp = malloc(rawExifLen);  // Freed in JpegEncode, Case 5
+        char *pp = malloc(rawExifLen);  // Freed in JpegEncode, Case 6
         if (!pp) {
+            if (comment) {
+                free(comment);
+            }
             if (extra) {
                 free(extra);
             }
@@ -1134,6 +1156,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)->comment = comment;
+    ((JPEGENCODERSTATE *)encoder->state.context)->comment_size = comment_size;
     ((JPEGENCODERSTATE *)encoder->state.context)->extra = extra;
     ((JPEGENCODERSTATE *)encoder->state.context)->extra_size = extra_size;
     ((JPEGENCODERSTATE *)encoder->state.context)->rawExif = rawExif;

src/libImaging/Jpeg.h

@@ -92,6 +92,10 @@ typedef struct {
     /* in factors of DCTSIZE2 */
     int qtablesLen;
 
+    /* Comment */
+    char *comment;
+    size_t comment_size;
+
     /* Extra data (to be injected after header) */
     char *extra;
     int extra_size;

src/libImaging/JpegEncode.c

@@ -277,6 +277,13 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             }
 
         case 4:
+
+            if (context->comment) {
+                jpeg_write_marker(&context->cinfo, JPEG_COM, (unsigned char *)context->comment, context->comment_size);
+            }
+            state->state++;
+
+        case 5:
             if (1024 > context->destination.pub.free_in_buffer) {
                 break;
             }
@@ -301,7 +308,7 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             state->state++;
             /* fall through */
 
-        case 5:
+        case 6:
 
             /* Finish compression */
             if (context->destination.pub.free_in_buffer < 100) {
@@ -310,6 +317,10 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             jpeg_finish_compress(&context->cinfo);
 
             /* Clean up */
+            if (context->comment) {
+                free(context->comment);
+                context->comment = NULL;
+            }
             if (context->extra) {
                 free(context->extra);
                 context->extra = NULL;