diff --git a/src/blosc2/__init__.py b/src/blosc2/__init__.py index 6b910593..c9c3ab36 100644 --- a/src/blosc2/__init__.py +++ b/src/blosc2/__init__.py @@ -9,6 +9,8 @@ # ruff: noqa: E402 - Module level import not at top of file # ruff: noqa: F401 - `var` imported but unused +import contextlib +import os import platform from enum import Enum @@ -402,7 +404,14 @@ def _raise(exc): nthreads -= nthreads // 8 if not IS_WASM: # WASM does not support threading - numexpr.set_num_threads(nthreads) + # Only call set_num_threads if within NUMEXPR_MAX_THREADS limit to avoid warning + numexpr_max_env = os.environ.get("NUMEXPR_MAX_THREADS") + numexpr_max: int | None = None + if numexpr_max_env is not None: + with contextlib.suppress(ValueError): + numexpr_max = int(numexpr_max_env) + if numexpr_max is None or nthreads <= numexpr_max: + numexpr.set_num_threads(nthreads) # This import must be before ndarray and schunk from .storage import ( # noqa: I001 diff --git a/tests/test_numexpr_threads.py b/tests/test_numexpr_threads.py new file mode 100644 index 00000000..b256497f --- /dev/null +++ b/tests/test_numexpr_threads.py @@ -0,0 +1,38 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +import os +import subprocess +import sys + + +def test_numexpr_max_threads_no_warning(): + """Test that importing blosc2 with NUMEXPR_MAX_THREADS set does not produce a warning. + + When NUMEXPR_MAX_THREADS is set to a value lower than the number of threads + blosc2 would use, we should NOT call numexpr.set_num_threads() to avoid + the numexpr warning being printed to stderr. + """ + # Inherit the current environment but set NUMEXPR_MAX_THREADS to a low value + env = os.environ.copy() + env["NUMEXPR_MAX_THREADS"] = "1" + + result = subprocess.run( + [sys.executable, "-c", "import blosc2; print(blosc2.__version__)"], + capture_output=True, + text=True, + env=env, + check=True, + ) + + # Check that no warning about NUMEXPR_MAX_THREADS was printed + assert "NUMEXPR_MAX_THREADS" not in result.stderr, ( + f"Unexpected numexpr warning in stderr: {result.stderr}" + ) + assert "nthreads cannot be larger" not in result.stderr, ( + f"Unexpected numexpr warning in stderr: {result.stderr}" + )