Let users set max_threads for AvifDecoder (#50)

fdintino · Jul 3, 2024 · 9d7e5b0 · 9d7e5b0
1 parent df4f409
commit 9d7e5b0
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/src/pillow_avif/AvifImagePlugin.py b/src/pillow_avif/AvifImagePlugin.py
@@ -62,7 +62,7 @@ def load_seek(self, pos):
 
     def _open(self):
         self._decoder = _avif.AvifDecoder(
-            self.fp.read(), DECODE_CODEC_CHOICE, CHROMA_UPSAMPLING
+            self.fp.read(), DECODE_CODEC_CHOICE, CHROMA_UPSAMPLING, DEFAULT_MAX_THREADS
         )
 
         # Get info from decoder

diff --git a/src/pillow_avif/_avif.c b/src/pillow_avif/_avif.c
@@ -731,10 +731,11 @@ AvifDecoderNew(PyObject *self_, PyObject *args) {
     char *codec_str;
     avifCodecChoice codec;
     avifChromaUpsampling upsampling;
+    int max_threads = 0;
 
     avifResult result;
 
-    if (!PyArg_ParseTuple(args, "Sss", &avif_bytes, &codec_str, &upsampling_str)) {
+    if (!PyArg_ParseTuple(args, "Sssi", &avif_bytes, &codec_str, &upsampling_str, &max_threads)) {
         return NULL;
     }
 
@@ -782,10 +783,13 @@ AvifDecoderNew(PyObject *self_, PyObject *args) {
 
     self->decoder = avifDecoderCreate();
 #if AVIF_VERSION >= 80400
-    if (default_max_threads == 0) {
-        init_max_threads();
+    if (max_threads == 0) {
+        if (default_max_threads == 0) {
+            init_max_threads();
+        }
+        max_threads = default_max_threads;
     }
-    self->decoder->maxThreads = default_max_threads;
+    self->decoder->maxThreads = max_threads;
 #endif
     self->decoder->codecChoice = codec;