Allow saving bytes as comments

joshware · Mar 23, 2023 · 5eee4cf · 5eee4cf
1 parent d3923f7
commit 5eee4cf
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 17 deletions.
diff --git a/Tests/test_file_jpeg2k.py b/Tests/test_file_jpeg2k.py
@@ -373,13 +373,13 @@ def test_comment():
 
 
 def test_save_comment():
-    comment = "Created by Pillow"
-    out = BytesIO()
-    test_card.save(out, "JPEG2000", comment=comment)
-    out.seek(0)
+    for comment in ("Created by Pillow", b"Created by Pillow"):
+        out = BytesIO()
+        test_card.save(out, "JPEG2000", comment=comment)
+        out.seek(0)
 
-    with Image.open(out) as im:
-        assert im.info["comment"] == b"Created by Pillow"
+        with Image.open(out) as im:
+            assert im.info["comment"] == b"Created by Pillow"
 
     too_long_comment = " " * 65532
     with pytest.raises(ValueError):

diff --git a/src/PIL/Jpeg2KImagePlugin.py b/src/PIL/Jpeg2KImagePlugin.py
@@ -351,10 +351,12 @@ def _save(im, fp, filename):
     cinema_mode = info.get("cinema_mode", "no")
     mct = info.get("mct", 0)
     signed = info.get("signed", False)
-    fd = -1
     comment = info.get("comment")
+    if isinstance(comment, str):
+        comment = comment.encode()
     add_plt = info.get("add_plt", False)
 
+    fd = -1
     if hasattr(fp, "fileno"):
         try:
             fd = fp.fileno()

diff --git a/src/encode.c b/src/encode.c
@@ -1215,11 +1215,12 @@ PyImaging_Jpeg2KEncoderNew(PyObject *self, PyObject *args) {
     int sgnd = 0;
     Py_ssize_t fd = -1;
     char *comment = NULL;
+    Py_ssize_t comment_size;
     int add_plt = 0;
 
     if (!PyArg_ParseTuple(
             args,
-            "ss|OOOsOnOOOssbbnzp",
+            "ss|OOOsOnOOOssbbnz#p",
             &mode,
             &format,
             &offset,
@@ -1237,6 +1238,7 @@ PyImaging_Jpeg2KEncoderNew(PyObject *self, PyObject *args) {
             &sgnd,
             &fd,
             &comment,
+            &comment_size,
             &add_plt)) {
         return NULL;
     }
@@ -1319,25 +1321,23 @@ PyImaging_Jpeg2KEncoderNew(PyObject *self, PyObject *args) {
         }
     }
 
-    if (comment != NULL && strlen(comment) > 0) {
+    if (comment && comment_size > 0) {
         /* Size is stored as as an uint16, subtract 4 bytes for the header */
-        if (strlen(comment) >= 65531) {
+        if (comment_size >= 65531) {
             PyErr_SetString(
                 PyExc_ValueError,
                 "JPEG 2000 comment is too long");
             Py_DECREF(encoder);
             return NULL;
         }
 
-        context->comment = strdup(comment);
-
-        if (context->comment == NULL) {
-            PyErr_SetString(
-                PyExc_MemoryError,
-                "Couldn't allocate memory for JPEG 2000 comment");
+        char *p = malloc(comment_size);
+        if (!p) {
             Py_DECREF(encoder);
-            return NULL;
+            return ImagingError_MemoryError();
         }
+        memcpy(p, comment, comment_size + 1);
+        context->comment = p;
     }
 
     if (quality_layers && PySequence_Check(quality_layers)) {