Merge branch 'radarhere/jp2k_options' into jp2k_options

Tests/test_file_jpeg2k.py

@@ -372,6 +372,20 @@ def test_comment():
             pass
 
 
+def test_save_comment():
+    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"
+
+    too_long_comment = " " * 65531
+    with pytest.raises(ValueError):
+        test_card.save(out, "JPEG2000", comment=too_long_comment)
+
+
 @pytest.mark.parametrize(
     "test_file",
     [
@@ -391,20 +405,6 @@ def test_crashes(test_file):
                 pass
 
 
-def test_custom_comment():
-    output_stream = BytesIO()
-    unique_comment = "This is a unique comment, which should be found below"
-    test_card.save(output_stream, "JPEG2000", comment=unique_comment)
-    output_stream.seek(0)
-    data = output_stream.read()
-    # Lazy method to determine if the comment is in the image generated
-    assert bytes(unique_comment, "utf-8") in data
-
-    too_long_comment = " " * 65532
-    with pytest.raises(ValueError):
-        test_card.save(output_stream, "JPEG2000", comment=too_long_comment)
-
-
 @skip_unless_feature_version("jpg_2000", "2.4.0")
 def test_plt_marker():
     # Search the start of the codesteam for the PLT box (id 0xFF58)

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()

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,9 +1321,9 @@ 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");
@@ -1329,15 +1331,14 @@ PyImaging_Jpeg2KEncoderNew(PyObject *self, PyObject *args) {
             return NULL;
         }
 
-        context->comment = strdup(comment);
+        char *p = malloc(comment_size + 1);
-
+        if (!p) {
-        if (context->comment == NULL) {
-            PyErr_SetString(
-                PyExc_MemoryError,
-                "Couldn't allocate memory for JPEG 2000 comment");
             Py_DECREF(encoder);
-            return NULL;
+            return ImagingError_MemoryError();
         }
+        memcpy(p, comment, comment_size);
+        p[comment_size] = '\0';
+        context->comment = p;
     }
 
     if (quality_layers && PySequence_Check(quality_layers)) {