Skip to content

Commit

Permalink
Make lazy_load.load thread-safe
Browse files Browse the repository at this point in the history
Closes #88
  • Loading branch information
stefanv committed Jan 25, 2024
1 parent bf82b68 commit d56c717
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 36 deletions.
79 changes: 43 additions & 36 deletions lazy_loader/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,16 @@
import inspect
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.
Expand Down Expand Up @@ -171,42 +175,45 @@ def myfunc():
Actual loading of the module occurs upon first attribute request.
"""
try:
return sys.modules[fullname]
except KeyError:
pass

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 = importlib.util.find_spec(fullname)
if spec is None:
if error_on_import:
raise ModuleNotFoundError(f"No module named '{fullname}'")
else:
try:
parent = inspect.stack()[1]
frame_data = {
"spec": fullname,
"filename": parent.filename,
"lineno": parent.lineno,
"function": parent.function,
"code_context": parent.code_context,
}
return DelayedImportErrorModule(frame_data, "DelayedImportErrorModule")
finally:
del parent

module = importlib.util.module_from_spec(spec)
sys.modules[fullname] = module

loader = importlib.util.LazyLoader(spec.loader)
loader.exec_module(module)
with threadlock:
try:
return sys.modules[fullname]
except KeyError:
pass

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 = importlib.util.find_spec(fullname)
if spec is None:
if error_on_import:
raise ModuleNotFoundError(f"No module named '{fullname}'")
else:
try:
parent = inspect.stack()[1]
frame_data = {
"spec": fullname,
"filename": parent.filename,
"lineno": parent.lineno,
"function": parent.function,
"code_context": parent.code_context,
}
return DelayedImportErrorModule(
frame_data, "DelayedImportErrorModule"
)
finally:
del parent

module = importlib.util.module_from_spec(spec)
sys.modules[fullname] = module

loader = importlib.util.LazyLoader(spec.loader)
loader.exec_module(module)

return module

Expand Down
13 changes: 13 additions & 0 deletions lazy_loader/tests/import_np_parallel.py
Original file line number Diff line number Diff line change
@@ -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()
13 changes: 13 additions & 0 deletions lazy_loader/tests/test_lazy_loader.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import importlib
import os
import subprocess
import sys
import types

Expand Down Expand Up @@ -149,3 +151,14 @@ def test_stub_loading_errors(tmp_path):

with pytest.raises(ValueError, match="Cannot load imports from non-existent stub"):
lazy.attach_stub("name", "not a file")


def test_parallel_load():
pytest.importorskip("numpy")

subprocess.run(
[
sys.executable,
os.path.join(os.path.dirname(__file__), "import_np_parallel.py"),
]
)

0 comments on commit d56c717

Please sign in to comment.