diff --git a/addon/flamenco/wheels/__init__.py b/addon/flamenco/wheels/__init__.py index eb0c601d..5a9a2d88 100644 --- a/addon/flamenco/wheels/__init__.py +++ b/addon/flamenco/wheels/__init__.py @@ -3,25 +3,31 @@ """External dependencies loader.""" import contextlib +import importlib from pathlib import Path import sys import logging from types import ModuleType -from typing import Iterator, Optional +from typing import Iterator _my_dir = Path(__file__).parent _log = logging.getLogger(__name__) -def load_wheel(module_name: str, fname_prefix: str) -> ModuleType: +def load_wheel(module_name: str, fname_prefix: str = "") -> ModuleType: """Loads a wheel from 'fname_prefix*.whl', unless the named module can be imported. This allows us to use system-installed packages before falling back to the shipped wheels. This is useful for development, less so for deployment. + + If `fname_prefix` is the empty string, it will use the module name. """ + if not fname_prefix: + fname_prefix = _fname_prefix_from_module_name(module_name) + try: - module = __import__(module_name) + module = importlib.import_module(module_name) except ImportError as ex: _log.debug("Unable to import %s directly, will try wheel: %s", module_name, ex) else: @@ -42,7 +48,7 @@ def load_wheel(module_name: str, fname_prefix: str) -> ModuleType: # this add-on from other add-ons. with _sys_path_mod_backup(wheel): try: - module = __import__(module_name) + module = importlib.import_module(module_name) except ImportError as ex: raise ImportError( "Unable to load %r from %s: %s" % (module_name, wheel, ex) @@ -53,15 +59,21 @@ def load_wheel(module_name: str, fname_prefix: str) -> ModuleType: return module -def load_wheel_global(module_name: str, fname_prefix: str) -> ModuleType: +def load_wheel_global(module_name: str, fname_prefix: str = "") -> ModuleType: """Loads a wheel from 'fname_prefix*.whl', unless the named module can be imported. This allows us to use system-installed packages before falling back to the shipped wheels. This is useful for development, less so for deployment. + + If `fname_prefix` is the empty string, it will use the first package from `module_name`. + In other words, `module_name="pkg.subpkg"` will result in `fname_prefix="pkg"`. """ + if not fname_prefix: + fname_prefix = _fname_prefix_from_module_name(module_name) + try: - module = __import__(module_name) + module = importlib.import_module(module_name) except ImportError as ex: _log.debug("Unable to import %s directly, will try wheel: %s", module_name, ex) else: @@ -80,7 +92,7 @@ def load_wheel_global(module_name: str, fname_prefix: str) -> ModuleType: sys.path.insert(0, wheel_filepath) try: - module = __import__(module_name) + module = importlib.import_module(module_name) except ImportError as ex: raise ImportError( "Unable to load %r from %s: %s" % (module_name, wheel, ex) @@ -92,16 +104,24 @@ def load_wheel_global(module_name: str, fname_prefix: str) -> ModuleType: @contextlib.contextmanager def _sys_path_mod_backup(wheel_file: Path) -> Iterator[None]: + """Temporarily inserts a wheel onto sys.path. + + When the context exits, it restores sys.path and sys.modules, so that + anything that was imported within the context remains unimportable by other + modules. + """ old_syspath = sys.path[:] + old_sysmod = sys.modules.copy() try: sys.path.insert(0, str(wheel_file)) yield finally: - # Restore without assigning new instances. That way references held by - # other code will stay valid. - + # Restore without assigning a new list instance. That way references + # held by other code will stay valid. sys.path[:] = old_syspath + sys.modules.clear() + sys.modules.update(old_sysmod) def _wheel_filename(fname_prefix: str) -> Path: @@ -119,6 +139,10 @@ def _wheel_filename(fname_prefix: str) -> Path: return wheels[-1] +def _fname_prefix_from_module_name(module_name: str) -> str: + return module_name.split(".", 1)[0] + + if __name__ == "__main__": logging.basicConfig(level=logging.DEBUG) wheel = _wheel_filename("python_dateutil")