jpegio: split open & decode; add support for decoding from files; add test

Author: jepler

locale/circuitpython.pot

@@ -212,6 +212,10 @@ msgstr ""
 msgid "%q must be of type %q or %q, not %q"
 msgstr ""
 
+#: shared-bindings/jpegio/JpegDecoder.c
+msgid "%q must be of type %q, %q, or %q, not %q"
+msgstr ""
+
 #: py/argcheck.c shared-module/synthio/__init__.c
 msgid "%q must be of type %q, not %q"
 msgstr ""
@@ -249,6 +253,10 @@ msgstr ""
 msgid "%q() takes %d positional arguments but %d were given"
 msgstr ""
 
+#: shared-module/jpegio/JpegDecoder.c
+msgid "%q() without %q()"
+msgstr ""
+
 #: shared-bindings/usb_hid/Device.c
 msgid "%q, %q, and %q must all be the same length"
 msgstr ""

shared-bindings/jpegio/JpegDecoder.c

@@ -25,6 +25,7 @@
  */
 
 #include "py/obj.h"
+#include "py/builtin.h"
 #include "py/runtime.h"
 
 #include "shared-bindings/displayio/Bitmap.h"
@@ -35,8 +36,19 @@
 //| class JpegDecoder:
 //|     """A JPEG decoder
 //|
-//|     A JPEG decoder allocates a few thousand bytes of memory. To reduce memory fragmentation,
+//|     A JpegDecoder allocates a few thousand bytes of memory. To reduce memory fragmentation,
 //|     create a single JpegDecoder object and use it anytime a JPEG image needs to be decoded.
+//|
+//|     Example::
+//|
+//|         from jpegio import JpegDecoder
+//|         from displayio import Bitmap
+//|
+//|         decoder = JpegDecoder()
+//|         width, height = decoder.open("/sd/example.jpg")
+//|         bitmap = Bitmap(width, height)
+//|         decoder.decode(bitmap)
+//|         # .. do something with bitmap
 //|     """
 //|
 //|     def __init__(self) -> None: ...
@@ -54,51 +66,88 @@ STATIC mp_obj_t jpegio_jpegdecoder_make_new(const mp_obj_type_t *type, size_t n_
     return MP_OBJ_FROM_PTR(self);
 }
 
+//|     @overload
+//|     def open(self, filename: str) -> Tuple[int, int]: ...
+//|     @overload
+//|     def open(self, buffer: ReadableBuffer) -> Tuple[int, int]: ...
+//|     @overload
+//|     def open(self, bytesio: io.BytesIO) -> Tuple[int, int]:
+//|         """Use the specified object  as the JPEG data source.
+//|
+//|         The source may be a filename, a binary buffer in memory, or an opened binary stream.
+//|
+//|         The single parameter is positional-only (write ``open(f)``, not
+//|         ``open(filename=f)`` but due to technical limitations this is
+//|         not shown in the function signature in the documentation.
+//|
+//|         Returns the image size as the tuple ``(width, height)``."""
+STATIC mp_obj_t jpegio_jpegdecoder_open(mp_obj_t self_in, mp_obj_t arg) {
+    jpegio_jpegdecoder_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    if (mp_obj_is_str(arg)) {
+        arg = mp_call_function_2(
+            MP_OBJ_FROM_PTR(&mp_builtin_open_obj),
+            arg,
+            MP_OBJ_NEW_QSTR(MP_QSTR_rb));
+    }
+
+    mp_buffer_info_t bufinfo;
+    const mp_stream_p_t *proto = mp_get_stream(arg);
+
+    if (proto && proto->read && !proto->is_text) {
+        return common_hal_jpegio_jpegdecoder_set_source_file(self, arg);
+    } else if (mp_get_buffer(arg, &bufinfo, MP_BUFFER_READ)) {
+        return common_hal_jpegio_jpegdecoder_set_source_buffer(self, arg);
+    }
+    mp_raise_TypeError_varg(MP_ERROR_TEXT("%q must be of type %q, %q, or %q, not %q"), MP_QSTR_data_source, MP_QSTR_str, MP_QSTR_BytesIO, MP_QSTR_ReadableBuffer);
+}
+MP_DEFINE_CONST_FUN_OBJ_2(jpegio_jpegdecoder_open_obj, jpegio_jpegdecoder_open);
+
 //|     def decode(
-//|         self, data: ReadableBuffer, bitmap: displayio.Bitmap | None = None, scale=0
-//|     ) -> tuple[int, int]:
+//|         self,
+//|         bitmap: displayio.Bitmap,
+//|         scale: int = 0,
+//|     ) -> None:
 //|         """Decode JPEG data
 //|
-//|         If ``bitmap`` is None, only the header is decoded.
-//|         Otherwise, the bitmap must be large enough to contain the decoded image.
+//|         The bitmap must be large enough to contain the decoded image.
 //|         The pixel data is stored in the `displayio.Colorspace.RGB565_SWAPPED` colorspace.
 //|
 //|         The image is optionally downscaled by a factor of ``2**scale``.
 //|         Scaling by a factor of 8 (scale=3) is particularly efficient in terms of decoding time.
 //|
-//|         :param ReadableBuffer data: Data in JPEG format
+//|         After a call to ``decode``, you must ``open`` a new JPEG. It is not
+//|         possible to repeatedly ``decode`` the same jpeg data, even if it is to
+//|         select different scales or crop regions from it.
+//|
 //|         :param Bitmap bitmap: Optional output buffer
-//|         :param int scale: Scale factor from 0 to 3.
-//|         :returns: The size of the (possibly scaled) image as ``width, height``
+//|         :param int scale: Scale factor from 0 to 3, inclusive.
 //|         """
 //|
 STATIC mp_obj_t jpegio_jpegdecoder_decode(mp_uint_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
     jpegio_jpegdecoder_obj_t *self = MP_OBJ_TO_PTR(pos_args[0]);
 
-    enum { ARG_data, ARG_bitmap, ARG_scale };
+    enum { ARG_bitmap, ARG_scale };
     static const mp_arg_t allowed_args[] = {
-        { MP_QSTR_data, MP_ARG_OBJ | MP_ARG_REQUIRED, {.u_obj = mp_const_none } },
-        { MP_QSTR_bitmap, MP_ARG_OBJ, {.u_obj = mp_const_none } },
+        { MP_QSTR_bitmap, MP_ARG_OBJ | MP_ARG_REQUIRED, {.u_obj = mp_const_none } },
         { MP_QSTR_scale, MP_ARG_INT, {.u_int = 0 } },
     };
     mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
     mp_arg_parse_all(n_args - 1, pos_args + 1, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
 
-    mp_buffer_info_t bufinfo;
-    mp_get_buffer_raise(args[ARG_data].u_obj, &bufinfo, MP_BUFFER_READ);
-
     mp_obj_t bitmap_in = args[ARG_bitmap].u_obj;
-    mp_arg_validate_type_or_none(bitmap_in, &displayio_bitmap_type, MP_QSTR_bitmap);
-    displayio_bitmap_t *bitmap = (args[ARG_bitmap].u_obj != mp_const_none) ? MP_OBJ_TO_PTR(args[ARG_bitmap].u_obj) : NULL;
+    mp_arg_validate_type(bitmap_in, &displayio_bitmap_type, MP_QSTR_bitmap);
+    displayio_bitmap_t *bitmap = MP_OBJ_TO_PTR(args[ARG_bitmap].u_obj);
 
     int scale = args[ARG_scale].u_int;
     mp_arg_validate_int_range(scale, 0, 3, MP_QSTR_scale);
 
-    return common_hal_jpegio_jpegdecoder_decode(self, bitmap, &bufinfo, scale);
+    common_hal_jpegio_jpegdecoder_decode_into(self, bitmap, scale);
+    return mp_const_none;
 }
-STATIC MP_DEFINE_CONST_FUN_OBJ_KW(jpegio_jpegdecoder_decode_obj, 2, jpegio_jpegdecoder_decode);
+STATIC MP_DEFINE_CONST_FUN_OBJ_KW(jpegio_jpegdecoder_decode_obj, 1, jpegio_jpegdecoder_decode);
 
 STATIC const mp_rom_map_elem_t jpegio_jpegdecoder_locals_dict_table[] = {
+    { MP_ROM_QSTR(MP_QSTR_open), MP_ROM_PTR(&jpegio_jpegdecoder_open_obj) },
     { MP_ROM_QSTR(MP_QSTR_decode), MP_ROM_PTR(&jpegio_jpegdecoder_decode_obj) },
 };
 STATIC MP_DEFINE_CONST_DICT(jpegio_jpegdecoder_locals_dict, jpegio_jpegdecoder_locals_dict_table);

shared-bindings/jpegio/JpegDecoder.h

@@ -25,11 +25,15 @@
  */
 
 #include "py/obj.h"
+#include "py/stream.h"
 #include "shared-module/displayio/Bitmap.h"
 
 extern const mp_obj_type_t jpegio_jpegdecoder_type;
 
 typedef struct jpegio_jpegdecoder_obj jpegio_jpegdecoder_obj_t;
 
 void common_hal_jpegio_jpegdecoder_construct(jpegio_jpegdecoder_obj_t *self);
-mp_obj_t common_hal_jpegio_jpegdecoder_decode(jpegio_jpegdecoder_obj_t *self, displayio_bitmap_t *bitmap, mp_buffer_info_t *jpeg_data, int scale);
+void common_hal_jpegio_jpegdecoder_close(jpegio_jpegdecoder_obj_t *self);
+mp_obj_t common_hal_jpegio_jpegdecoder_set_source_buffer(jpegio_jpegdecoder_obj_t *self, mp_obj_t jpeg_data);
+mp_obj_t common_hal_jpegio_jpegdecoder_set_source_file(jpegio_jpegdecoder_obj_t *self, mp_obj_t file_obj);
+void common_hal_jpegio_jpegdecoder_decode_into(jpegio_jpegdecoder_obj_t *self, displayio_bitmap_t *bitmap, int scale);

shared-module/jpegio/JpegDecoder.c

@@ -30,6 +30,8 @@
 #include "shared-bindings/bitmaptools/__init__.h"
 #include "shared-module/jpegio/JpegDecoder.h"
 
+typedef size_t (*input_func)(JDEC *jd, uint8_t *dest, size_t len);
+
 // Given a pointer `ptr` to the field `field_name` inside a structure of type `type`,
 // retrieve a pointer to the containing object.
 // This is used to retrieve the jpegio_jpegdecoder_obj_t object given the JDEC.
@@ -71,12 +73,61 @@ static void check_jresult(JRESULT j) {
 }
 
 void common_hal_jpegio_jpegdecoder_construct(jpegio_jpegdecoder_obj_t *self) {
-    // Nothing(!)
+    self->data_obj = MP_OBJ_NULL;
+}
+
+void common_hal_jpegio_jpegdecoder_close(jpegio_jpegdecoder_obj_t *self) {
+    self->data_obj = MP_OBJ_NULL;
+    memset(&self->bufinfo, 0, sizeof(self->bufinfo));
+}
+
+static mp_obj_t common_hal_jpegio_jpegdecoder_decode_common(jpegio_jpegdecoder_obj_t *self, input_func fun) {
+    JRESULT result = jd_prepare(&self->decoder, fun, &self->workspace, sizeof(self->workspace), NULL);
+    if (result != JDR_OK) {
+        common_hal_jpegio_jpegdecoder_close(self);
+    }
+    check_jresult(result);
+    mp_obj_t elems[] = {
+        MP_OBJ_NEW_SMALL_INT(self->decoder.height),
+        MP_OBJ_NEW_SMALL_INT(self->decoder.width)
+    };
+    return mp_obj_new_tuple(MP_ARRAY_SIZE(elems), elems);
+}
+
+static size_t file_input(JDEC *jd, uint8_t *dest, size_t len) {
+    jpegio_jpegdecoder_obj_t *self = CONTAINER_OF(jd, jpegio_jpegdecoder_obj_t, decoder);
+
+    if (!dest) {
+        // caller passes NULL to skip data; we need to read over the data.
+        // Don't assume a seekable stream, because we want to decode jpegs
+        // right from a native socket object
+        uint8_t buf[512];
+        size_t total = 0;
+        size_t read;
+        do {
+            size_t to_discard = MIN(len - total, sizeof(buf));
+            read = file_input(jd, buf, to_discard);
+            total += read;
+        } while (read != 0 && total != len);
+        return len;
+    }
+
+    int errcode = 0;
+    size_t result = mp_stream_rw(self->data_obj, dest, len, &errcode, MP_STREAM_RW_READ);
+    if (errcode != 0) { // raise our own error in case of I/O failure, it's better than the decoder's error
+        mp_raise_OSError(errcode);
+    }
+    return result;
+}
+
+mp_obj_t common_hal_jpegio_jpegdecoder_set_source_file(jpegio_jpegdecoder_obj_t *self, mp_obj_t file_obj) {
+    self->data_obj = file_obj;
+    return common_hal_jpegio_jpegdecoder_decode_common(self, file_input);
 }
 
 static size_t buffer_input(JDEC *jd, uint8_t *dest, size_t len) {
     jpegio_jpegdecoder_obj_t *self = CONTAINER_OF(jd, jpegio_jpegdecoder_obj_t, decoder);
-    mp_buffer_info_t *src = &self->src;
+    mp_buffer_info_t *src = &self->bufinfo;
     size_t to_copy = MIN(len, src->len);
     if (dest) { // passes NULL to skip data
         memcpy(dest, src->buf, to_copy);
@@ -86,24 +137,29 @@ static size_t buffer_input(JDEC *jd, uint8_t *dest, size_t len) {
     return to_copy;
 }
 
+mp_obj_t common_hal_jpegio_jpegdecoder_set_source_buffer(jpegio_jpegdecoder_obj_t *self, mp_obj_t buffer_obj) {
+    self->data_obj = buffer_obj;
+    mp_get_buffer_raise(buffer_obj, &self->bufinfo, MP_BUFFER_READ);
+    return common_hal_jpegio_jpegdecoder_decode_common(self, buffer_input);
+}
+
 static int bitmap_output(JDEC *jd, void *data, JRECT *rect) {
     jpegio_jpegdecoder_obj_t *self = CONTAINER_OF(jd, jpegio_jpegdecoder_obj_t, decoder);
     common_hal_bitmaptools_arrayblit(self->dest, data, 2, rect->left, rect->top, rect->right + 1, rect->bottom + 1, false, 0);
     return 1;
 }
 
-mp_obj_t common_hal_jpegio_jpegdecoder_decode(jpegio_jpegdecoder_obj_t *self, displayio_bitmap_t *bitmap, mp_buffer_info_t *jpeg_data, int scale) {
-    self->src = *jpeg_data;
-    check_jresult(jd_prepare(&self->decoder, buffer_input, &self->workspace, sizeof(self->workspace), NULL));
+void common_hal_jpegio_jpegdecoder_decode_into(jpegio_jpegdecoder_obj_t *self, displayio_bitmap_t *bitmap, int scale) {
+    if (self->data_obj == MP_OBJ_NULL) {
+        mp_raise_RuntimeError_varg(MP_ERROR_TEXT("%q() without %q()"), MP_QSTR_decode, MP_QSTR_open);
+    }
     int dst_height = self->decoder.height >> scale;
     int dst_width = self->decoder.width >> scale;
-    if (bitmap) {
-        if (dst_width > bitmap->width || dst_height > bitmap->height) {
-            mp_raise_ValueError(MP_ERROR_TEXT("Destination bitmap too small to contain image"));
-        }
-        self->dest = bitmap;
-        check_jresult(jd_decomp(&self->decoder, bitmap_output, scale));
+    if (dst_width > bitmap->width || dst_height > bitmap->height) {
+        mp_raise_ValueError(MP_ERROR_TEXT("Destination bitmap too small to contain image"));
     }
-    mp_obj_t elems[] = { MP_OBJ_NEW_SMALL_INT(dst_width), MP_OBJ_NEW_SMALL_INT(dst_height) };
-    return mp_obj_new_tuple(MP_ARRAY_SIZE(elems), elems);
+    self->dest = bitmap;
+    JRESULT result = jd_decomp(&self->decoder, bitmap_output, scale);
+    common_hal_jpegio_jpegdecoder_close(self);
+    check_jresult(result);
 }

shared-module/jpegio/JpegDecoder.h

@@ -34,6 +34,8 @@ typedef struct jpegio_jpegdecoder_obj {
     mp_obj_base_t base;
     JDEC decoder;
     byte workspace[TJPGD_WORKSPACE_SIZE];
-    mp_buffer_info_t src;
+    mp_obj_t data_obj;
+    mp_buffer_info_t bufinfo;
     displayio_bitmap_t *dest;
+    int scale;
 } jpegio_jpegdecoder_obj_t;

tests/circuitpython/jpegio_decompress.py

@@ -1,3 +1,5 @@
+import io
+
 from displayio import Bitmap
 import binascii
 import jpegio
@@ -53,12 +55,14 @@
 
 
 def test(jpeg_input, scale):
-    w, h = decoder.decode(jpeg_input, scale=scale)
+    w, h = decoder.open(jpeg_input)
+    w >>= scale
+    h >>= scale
     print(f"{w}x{h}")
 
     b = Bitmap(w, h, 65535)
 
-    decoder.decode(jpeg_input, b, scale=scale)
+    decoder.decode(b, scale=scale)
 
     for y in range(min(30, h)):
         for x in range(min(72, w)):
@@ -67,7 +71,16 @@ def test(jpeg_input, scale):
     print("...\n" if h > 24 else "")
 
 
+print("bytes")
 test(content, scale=0)
 test(content, scale=1)
 test(content, scale=2)
 test(content, scale=3)
+
+bytesio = io.BytesIO(content)
+
+print("BytesIO")
+test(io.BytesIO(content), scale=0)
+test(io.BytesIO(content), scale=1)
+test(io.BytesIO(content), scale=2)
+test(io.BytesIO(content), scale=3)