From a33f3f9fe9cc753cf125592e0daf6ef624887fda Mon Sep 17 00:00:00 2001 From: Stefan van der Walt Date: Tue, 30 Jan 2024 14:39:45 -0800 Subject: [PATCH] Make `lazy_load.load` partially thread-safe (#90) --- lazy_loader/__init__.py | 125 +++++++++++++----------- lazy_loader/tests/import_np_parallel.py | 13 +++ lazy_loader/tests/test_lazy_loader.py | 13 +++ 3 files changed, 92 insertions(+), 59 deletions(-) create mode 100644 lazy_loader/tests/import_np_parallel.py diff --git a/lazy_loader/__init__.py b/lazy_loader/__init__.py index f0f4672..6bec8f8 100644 --- a/lazy_loader/__init__.py +++ b/lazy_loader/__init__.py @@ -9,12 +9,16 @@ import importlib.util import os import sys +import threading import types import warnings __all__ = ["attach", "load", "attach_stub"] +threadlock = threading.Lock() + + def attach(package_name, submodules=None, submod_attrs=None): """Attach lazily loaded submodules, functions, or other attributes. @@ -179,66 +183,69 @@ def myfunc(): Actual loading of the module occurs upon first attribute request. """ - module = sys.modules.get(fullname) - have_module = module is not None - - # Most common, short-circuit - if have_module and require is None: - return module - - if "." in fullname: - msg = ( - "subpackages can technically be lazily loaded, but it causes the " - "package to be eagerly loaded even if it is already lazily loaded." - "So, you probably shouldn't use subpackages with this lazy feature." - ) - warnings.warn(msg, RuntimeWarning) - - spec = None - if not have_module: - spec = importlib.util.find_spec(fullname) - have_module = spec is not None - - if not have_module: - not_found_message = f"No module named '{fullname}'" - elif require is not None: - try: - have_module = _check_requirement(require) - except ModuleNotFoundError as e: - raise ValueError( - f"Found module '{fullname}' but cannot test requirement '{require}'. " - "Requirements must match distribution name, not module name." - ) from e - - not_found_message = f"No distribution can be found matching '{require}'" - - if not have_module: - if error_on_import: - raise ModuleNotFoundError(not_found_message) - import inspect - - try: - parent = inspect.stack()[1] - frame_data = { - "filename": parent.filename, - "lineno": parent.lineno, - "function": parent.function, - "code_context": parent.code_context, - } - return DelayedImportErrorModule( - frame_data, - "DelayedImportErrorModule", - message=not_found_message, + with threadlock: + module = sys.modules.get(fullname) + have_module = module is not None + + # Most common, short-circuit + if have_module and require is None: + return module + + if "." in fullname: + msg = ( + "subpackages can technically be lazily loaded, but it causes the " + "package to be eagerly loaded even if it is already lazily loaded." + "So, you probably shouldn't use subpackages with this lazy feature." ) - finally: - del parent - - if spec is not None: - module = importlib.util.module_from_spec(spec) - sys.modules[fullname] = module - - loader = importlib.util.LazyLoader(spec.loader) - loader.exec_module(module) + warnings.warn(msg, RuntimeWarning) + + spec = None + + if not have_module: + spec = importlib.util.find_spec(fullname) + have_module = spec is not None + + if not have_module: + not_found_message = f"No module named '{fullname}'" + elif require is not None: + try: + have_module = _check_requirement(require) + except ModuleNotFoundError as e: + raise ValueError( + f"Found module '{fullname}' but cannot test " + "requirement '{require}'. " + "Requirements must match distribution name, not module name." + ) from e + + not_found_message = f"No distribution can be found matching '{require}'" + + if not have_module: + if error_on_import: + raise ModuleNotFoundError(not_found_message) + import inspect + + try: + parent = inspect.stack()[1] + frame_data = { + "filename": parent.filename, + "lineno": parent.lineno, + "function": parent.function, + "code_context": parent.code_context, + } + return DelayedImportErrorModule( + frame_data, + "DelayedImportErrorModule", + message=not_found_message, + ) + finally: + del parent + + if spec is not None: + module = importlib.util.module_from_spec(spec) + sys.modules[fullname] = module + + loader = importlib.util.LazyLoader(spec.loader) + loader.exec_module(module) return module diff --git a/lazy_loader/tests/import_np_parallel.py b/lazy_loader/tests/import_np_parallel.py new file mode 100644 index 0000000..50dc0d0 --- /dev/null +++ b/lazy_loader/tests/import_np_parallel.py @@ -0,0 +1,13 @@ +import threading +import time + +import lazy_loader as lazy + + +def import_np(): + time.sleep(0.5) + lazy.load("numpy") + + +for _ in range(10): + threading.Thread(target=import_np).start() diff --git a/lazy_loader/tests/test_lazy_loader.py b/lazy_loader/tests/test_lazy_loader.py index a7c166c..f28439f 100644 --- a/lazy_loader/tests/test_lazy_loader.py +++ b/lazy_loader/tests/test_lazy_loader.py @@ -1,4 +1,6 @@ import importlib +import os +import subprocess import sys import types from unittest import mock @@ -172,3 +174,14 @@ def test_require_kwarg(): # raise a ValueError with pytest.raises(ValueError): lazy.load("math", require="somepkg >= 1.0") + + +def test_parallel_load(): + pytest.importorskip("numpy") + + subprocess.run( + [ + sys.executable, + os.path.join(os.path.dirname(__file__), "import_np_parallel.py"), + ] + )