Set default max threads in Python (#4)

* Removed unused C values * Set default max threads in Python --------- Co-authored-by: Andrew Murray <radarhere@users.noreply.github.com>
fdintino · Nov 11, 2024 · 50b993a · 50b993a
1 parent d6a0a15
commit 50b993a
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 77 deletions.
diff --git a/src/PIL/AvifImagePlugin.py b/src/PIL/AvifImagePlugin.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from io import BytesIO
 from typing import IO
 
@@ -46,6 +47,15 @@ def _accept(prefix: bytes) -> bool | str:
     return False
 
 
+def _get_default_max_threads():
+    if DEFAULT_MAX_THREADS:
+        return DEFAULT_MAX_THREADS
+    if hasattr(os, "sched_getaffinity"):
+        return len(os.sched_getaffinity(0))
+    else:
+        return os.cpu_count() or 1
+
+
 class AvifImageFile(ImageFile.ImageFile):
     format = "AVIF"
     format_description = "AVIF image"
@@ -64,7 +74,10 @@ def _open(self) -> None:
             raise SyntaxError(msg)
 
         self._decoder = _avif.AvifDecoder(
-            self.fp.read(), DECODE_CODEC_CHOICE, CHROMA_UPSAMPLING, DEFAULT_MAX_THREADS
+            self.fp.read(),
+            DECODE_CODEC_CHOICE,
+            CHROMA_UPSAMPLING,
+            _get_default_max_threads(),
         )
 
         # Get info from decoder
@@ -140,7 +153,7 @@ def _save(
     duration = info.get("duration", 0)
     subsampling = info.get("subsampling", "4:2:0")
     speed = info.get("speed", 6)
-    max_threads = info.get("max_threads", DEFAULT_MAX_THREADS)
+    max_threads = info.get("max_threads", _get_default_max_threads())
     codec = info.get("codec", "auto")
     range_ = info.get("range", "full")
     tile_rows_log2 = info.get("tile_rows", 0)

diff --git a/src/_avif.c b/src/_avif.c
@@ -44,54 +44,6 @@ typedef struct {
 
 static PyTypeObject AvifDecoder_Type;
 
-static int default_max_threads = 0;
-
-static void
-init_max_threads(void) {
-    PyObject *os = NULL;
-    PyObject *n = NULL;
-    long num_cpus;
-
-    os = PyImport_ImportModule("os");
-    if (os == NULL) {
-        goto error;
-    }
-
-    if (PyObject_HasAttrString(os, "sched_getaffinity")) {
-        n = PyObject_CallMethod(os, "sched_getaffinity", "i", 0);
-        if (n == NULL) {
-            goto error;
-        }
-        num_cpus = PySet_Size(n);
-    } else {
-        n = PyObject_CallMethod(os, "cpu_count", NULL);
-        if (n == NULL) {
-            goto error;
-        }
-        num_cpus = PyLong_AsLong(n);
-    }
-
-    if (num_cpus < 1) {
-        goto error;
-    }
-
-    default_max_threads = (int)num_cpus;
-
-done:
-    Py_XDECREF(os);
-    Py_XDECREF(n);
-    return;
-
-error:
-    if (PyErr_Occurred()) {
-        PyErr_Clear();
-    }
-    PyErr_WarnEx(
-        PyExc_RuntimeWarning, "could not get cpu count: using max_threads=1", 1
-    );
-    goto done;
-}
-
 static int
 normalize_quantize_value(int qvalue) {
     if (qvalue < AVIF_QUANTIZER_BEST_QUALITY) {
@@ -306,23 +258,23 @@ AvifEncoderNew(PyObject *self_, PyObject *args) {
     AvifEncoderObject *self = NULL;
     avifEncoder *encoder = NULL;
 
-    char *subsampling = "4:2:0";
-    int qmin = AVIF_QUANTIZER_BEST_QUALITY;  // =0
-    int qmax = 10;                           // "High Quality", but not lossless
-    int quality = 75;
-    int speed = 8;
-    int exif_orientation = 0;
-    int max_threads = default_max_threads;
+    char *subsampling;
+    int qmin;
+    int qmax;
+    int quality;
+    int speed;
+    int exif_orientation;
+    int max_threads;
     PyObject *icc_bytes;
     PyObject *exif_bytes;
     PyObject *xmp_bytes;
-    PyObject *alpha_premultiplied = NULL;
-    PyObject *autotiling = NULL;
-    int tile_rows_log2 = 0;
-    int tile_cols_log2 = 0;
+    PyObject *alpha_premultiplied;
+    PyObject *autotiling;
+    int tile_rows_log2;
+    int tile_cols_log2;
 
-    char *codec = "auto";
-    char *range = "full";
+    char *codec;
+    char *range;
 
     PyObject *advanced;
 
@@ -438,13 +390,6 @@ AvifEncoderNew(PyObject *self_, PyObject *args) {
 
         encoder = avifEncoderCreate();
 
-        if (max_threads == 0) {
-            if (default_max_threads == 0) {
-                init_max_threads();
-            }
-            max_threads = default_max_threads;
-        }
-
         int is_aom_encode = strcmp(codec, "aom") == 0 ||
                             (strcmp(codec, "auto") == 0 &&
                              _codec_available("aom", AVIF_CODEC_FLAG_CAN_ENCODE));
@@ -730,7 +675,7 @@ AvifDecoderNew(PyObject *self_, PyObject *args) {
     char *codec_str;
     avifCodecChoice codec;
     avifChromaUpsampling upsampling;
-    int max_threads = 0;
+    int max_threads;
 
     avifResult result;
 
@@ -785,12 +730,6 @@ AvifDecoderNew(PyObject *self_, PyObject *args) {
 
     self->decoder = avifDecoderCreate();
 #if AVIF_VERSION >= 80400
-    if (max_threads == 0) {
-        if (default_max_threads == 0) {
-            init_max_threads();
-        }
-        max_threads = default_max_threads;
-    }
     self->decoder->maxThreads = max_threads;
 #endif
 #if AVIF_VERSION >= 90200