Add-on: load all BAT submodules in one go

Adjust the loading of BAT from a wheel file in such a way that all
submodules are loaded in one go. This ensures that they're still
isolated from the rest of Blender (so other add-ons won't find our BAT),
but not from each other (so that there is only one copy of each
submodule).

In practice, this solves an issue where calling
`blender_asset_tracer.blendfile.set_strict_pointer_mode(False)` had no
effect. This was caused by each loaded submodule having a different copy
of `blendfile`.

Also loaded modules are logged more explicitly (at INFO level) to aid in
debugging later on.

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 = {}

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

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

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: