Skip to content

Commit

Permalink
Make lazy_load.load partially thread-safe (#90)
Browse files Browse the repository at this point in the history
  • Loading branch information
stefanv committed Jan 30, 2024
1 parent 0d4fb38 commit a33f3f9
Show file tree
Hide file tree
Showing 3 changed files with 92 additions and 59 deletions.
125 changes: 66 additions & 59 deletions lazy_loader/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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

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
from unittest import mock
Expand Down Expand Up @@ -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"),
]
)

0 comments on commit a33f3f9

Please sign in to comment.