diff --git a/tests/test_cache.py b/tests/test_cache.py index 90ff0c3a..ba86a905 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -1,5 +1,8 @@ """Tests for ``threadpoolctl``'s internal caching layer.""" +import os +import subprocess +import sys from ctypes import CDLL import pytest @@ -57,3 +60,52 @@ def cleanup(): # Only one call should ever happen when using the cached method, however: assert libs[0].my_extra_method() is initial assert libs[0].my_extra_method() is initial + + +@pytest.mark.skipif( + not sys.platform.startswith("linux"), reason="Relies on glibc being loaded" +) +def test_register_after_libraries_were_listed(): + """ + A controller registered after a library was already listed (and cached as + unsupported) still finds it. + """ + # Run in a subprocess to keep the registered controller out of the other + # tests: + path = os.path.dirname(os.path.dirname(__file__)) + env = os.environ.copy() + env["PYTHONPATH"] = path + os.pathsep + env.get("PYTHONPATH", "") + script = """ +from threadpoolctl import ( + LibController, + ThreadpoolController, + get_cached_controller, + register, +) + + +class FakeController(LibController): + user_api = "fake" + internal_api = "fake" + filename_prefixes = ("libc.",) + check_symbols = ("printf",) + + def get_num_threads(self): + return 1 + + def set_num_threads(self, num_threads): + pass + + def get_version(self): + return None + + +# Lists glibc, which no controller matches yet: +assert not ThreadpoolController().select(user_api="fake").lib_controllers +assert not get_cached_controller().select(user_api="fake").lib_controllers + +register(FakeController) +assert ThreadpoolController().select(user_api="fake").lib_controllers +assert get_cached_controller().select(user_api="fake").lib_controllers +""" + subprocess.check_call([sys.executable, "-c", script], env=env) diff --git a/threadpoolctl.py b/threadpoolctl.py index 47bc8e15..34e4b9b2 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -975,6 +975,9 @@ def register(controller): _ALL_USER_APIS.append(controller.user_api) _ALL_INTERNAL_APIS.append(controller.internal_api) _ALL_PREFIXES.extend(controller.filename_prefixes) + # Libraries listed before this call may match the new controller: + _CDLL_CACHE._controller_cache.clear() + reset_cached_controller() def _format_docstring(*args, **kwargs):