diff --git a/addon/flamenco/bat/interface.py b/addon/flamenco/bat/interface.py index 9662fab3..07b50880 100644 --- a/addon/flamenco/bat/interface.py +++ b/addon/flamenco/bat/interface.py @@ -9,13 +9,7 @@ import queue import threading import typing -from .. import wheels - -blendfile = wheels.load_wheel("blender_asset_tracer.blendfile") -pack = wheels.load_wheel("blender_asset_tracer.pack") -progress = wheels.load_wheel("blender_asset_tracer.pack.progress") -transfer = wheels.load_wheel("blender_asset_tracer.pack.transfer") -shaman = wheels.load_wheel("blender_asset_tracer.pack.shaman") +from . import submodules log = logging.getLogger(__name__) @@ -58,7 +52,7 @@ class MsgDone(Message): # MyPy doesn't understand the way BAT subpackages are imported. -class BatProgress(progress.Callback): # type: ignore +class BatProgress(submodules.progress.Callback): # type: ignore """Report progress of BAT Packing to the given queue.""" def __init__(self, queue: queue.SimpleQueue[Message]) -> None: @@ -128,7 +122,7 @@ class PackThread(threading.Thread): queue: queue.SimpleQueue[Message] # MyPy doesn't understand the way BAT subpackages are imported. - def __init__(self, packer: pack.Packer) -> None: # type: ignore + def __init__(self, packer: submodules.pack.Packer) -> None: # type: ignore # Quitting Blender should abort the transfer (instead of hanging until # the transfer is done), hence daemon=True. super().__init__(daemon=True, name="PackThread") @@ -197,7 +191,7 @@ def copy( # type: ignore exclusion_filter: str, *, relative_only: bool, - packer_class=pack.Packer, + packer_class=submodules.pack.Packer, packer_kwargs: Optional[dict[Any, Any]] = None, ) -> PackThread: """Use BAT to copy the given file and dependencies to the target location. @@ -214,7 +208,7 @@ def copy( # type: ignore # Due to issues with library overrides and unsynced pointers, it's quite # common for the Blender Animation Studio to get crashes of BAT. To avoid # these, Strict Pointer Mode is disabled. - blendfile.set_strict_pointer_mode(False) + submodules.blendfile.set_strict_pointer_mode(False) if packer_kwargs is None: packer_kwargs = {} diff --git a/addon/flamenco/bat/shaman.py b/addon/flamenco/bat/shaman.py index 47b61155..2d3ca1f6 100644 --- a/addon/flamenco/bat/shaman.py +++ b/addon/flamenco/bat/shaman.py @@ -7,13 +7,7 @@ from collections import deque from pathlib import Path, PurePath, PurePosixPath from typing import TYPE_CHECKING, Optional, Any, Iterable, Iterator -from .. import wheels -from . import cache - -bat_pack = wheels.load_wheel("blender_asset_tracer.pack") -bat_transfer = wheels.load_wheel("blender_asset_tracer.pack.transfer") -bat_bpathlib = wheels.load_wheel("blender_asset_tracer.bpathlib") - +from . import cache, submodules if TYPE_CHECKING: from ..manager import ApiClient as _ApiClient @@ -37,8 +31,8 @@ MAX_FAILED_PATHS = 8 HashableShamanFileSpec = tuple[str, int, str] """Tuple of the 'sha', 'size', and 'path' fields of a ShamanFileSpec.""" -# Mypy doesn't understand that bat_pack.Packer exists. -class Packer(bat_pack.Packer): # type: ignore +# Mypy doesn't understand that submodules.pack.Packer exists. +class Packer(submodules.pack.Packer): # type: ignore """Creates BAT Packs on a Shaman server.""" def __init__( @@ -60,8 +54,8 @@ class Packer(bat_pack.Packer): # type: ignore self.api_client = api_client self.shaman_transferrer: Optional[Transferrer] = None - # Mypy doesn't understand that bat_transfer.FileTransferer exists. - def _create_file_transferer(self) -> bat_transfer.FileTransferer: # type: ignore + # Mypy doesn't understand that submodules.transfer.FileTransferer exists. + def _create_file_transferer(self) -> submodules.transfer.FileTransferer: # type: ignore self.shaman_transferrer = Transferrer( self.api_client, self.project, self.checkout_path ) @@ -90,7 +84,7 @@ class Packer(bat_pack.Packer): # type: ignore self._check_aborted() -class Transferrer(bat_transfer.FileTransferer): # type: ignore +class Transferrer(submodules.transfer.FileTransferer): # type: ignore """Sends files to a Shaman server.""" class AbortUpload(Exception): @@ -213,7 +207,7 @@ class Transferrer(bat_transfer.FileTransferer): # type: ignore try: checksum = cache.compute_cached_checksum(src) filesize = src.stat().st_size - relpath = str(bat_bpathlib.strip_root(dst)) + relpath = str(submodules.bpathlib.strip_root(dst)) filespec = ShamanFileSpec( sha=checksum, @@ -223,7 +217,7 @@ class Transferrer(bat_transfer.FileTransferer): # type: ignore filespecs.append(filespec) self._rel_to_local_path[relpath] = src - if act == bat_transfer.Action.MOVE: + if act == submodules.transfer.Action.MOVE: self._delete_when_done.append(src) except Exception: # We have to catch exceptions in a broad way, as this is running in diff --git a/addon/flamenco/bat/submodules.py b/addon/flamenco/bat/submodules.py new file mode 100644 index 00000000..399cdb26 --- /dev/null +++ b/addon/flamenco/bat/submodules.py @@ -0,0 +1,7 @@ +from .. import wheels + +# Load all the submodules we need from BAT in one go. +_bat_modules = wheels.load_wheel( + "blender_asset_tracer", + ("blendfile", "pack", "pack.progress", "pack.transfer", "pack.shaman", "bpathlib")) +bat_toplevel, blendfile, pack, progress, transfer, shaman, bpathlib = _bat_modules diff --git a/addon/flamenco/wheels/__init__.py b/addon/flamenco/wheels/__init__.py index 5a9a2d88..c7141397 100644 --- a/addon/flamenco/wheels/__init__.py +++ b/addon/flamenco/wheels/__init__.py @@ -13,50 +13,42 @@ from typing import Iterator _my_dir = Path(__file__).parent _log = logging.getLogger(__name__) +def load_wheel(module_name: str, submodules: tuple[str]) -> list[ModuleType]: + """Loads modules from a wheel file 'module_name*.whl'. -def load_wheel(module_name: str, fname_prefix: str = "") -> ModuleType: - """Loads a wheel from 'fname_prefix*.whl', unless the named module can be imported. + Loads `module_name`, and if submodules are given, loads + `module_name.submodule` for each of the submodules. This allows loading all + required modules from the same wheel in one session, ensuring that + inter-submodule references are correct. - 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. + Returns the loaded modules, so [module, submodule, submodule, ...]. """ - if not fname_prefix: - fname_prefix = _fname_prefix_from_module_name(module_name) - - try: - 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: - _log.debug( - "Was able to load %s from %s, no need to load wheel %s", - module_name, - module.__file__, - fname_prefix, - ) - assert isinstance(module, ModuleType) - return module - + fname_prefix = _fname_prefix_from_module_name(module_name) wheel = _wheel_filename(fname_prefix) + loaded_modules: list[ModuleType] = [] + to_load = [module_name] + [f"{module_name}.{submodule}" for submodule in submodules] + # Load the module from the wheel file. Keep a backup of sys.path so that it # can be restored later. This should ensure that future import statements # cannot find this wheel file, increasing the separation of dependencies of # this add-on from other add-ons. with _sys_path_mod_backup(wheel): - try: - module = importlib.import_module(module_name) - except ImportError as ex: - raise ImportError( - "Unable to load %r from %s: %s" % (module_name, wheel, ex) - ) from None + for modname in to_load: + try: + module = importlib.import_module(modname) + except ImportError as ex: + raise ImportError( + "Unable to load %r from %s: %s" % (modname, wheel, ex) + ) from None + assert isinstance(module, ModuleType) + loaded_modules.append(module) + _log.info("Loaded %s from %s", modname, module.__file__) - _log.debug("Loaded %s from %s", module_name, module.__file__) - assert isinstance(module, ModuleType) - return module + assert len(loaded_modules) == len(to_load), \ + f"expecting to load {len(to_load)} modules, but only have {len(loaded_modules)}: {loaded_modules}" + return loaded_modules def load_wheel_global(module_name: str, fname_prefix: str = "") -> ModuleType: