loader: add support for components

examples/loader_component.py

@@ -0,0 +1,11 @@
+# This example shows how you can use the `wasmtime.loader` module to load wasm
+# components without having to generate bindings manually
+
+import wasmtime, wasmtime.loader
+
+import loader_component_add  # type: ignore
+
+
+store = wasmtime.Store()
+component = loader_component_add.Root(store)
+assert component.add(store, 1, 2) == 3

examples/loader_component_add.wat

@@ -0,0 +1,12 @@
+(component
+  (core module $C
+    (func (export "add") (param i32 i32) (result i32)
+      local.get 0
+      local.get 1
+      i32.add)
+  )
+  (core instance $c (instantiate $C))
+  (core func $add (alias core export $c "add"))
+  (func (export "add") (param "x" s32) (param "y" s32) (result s32)
+    (canon lift (core func $add)))
+)

rust/src/bindgen.rs

@@ -539,16 +539,27 @@ impl<'a> Instantiator<'a> {
         let i = self.instances.push(idx);
         let core_file_name = self.gen.core_file_name(&self.name, idx.as_u32());
         self.gen.init.pyimport("pathlib", None);
+        self.gen.init.pyimport("importlib_resources", None);
 
         uwriteln!(
             self.gen.init,
-            "path = pathlib.Path(__file__).parent / ('{}')",
+            "file = importlib_resources.files() / ('{}')",
             core_file_name,
         );
+        uwriteln!(self.gen.init, "if isinstance(file, pathlib.Path):");
+        self.gen.init.indent();
         uwriteln!(
             self.gen.init,
-            "module = wasmtime.Module.from_file(store.engine, path)"
+            "module = wasmtime.Module.from_file(store.engine, file)"
         );
+        self.gen.init.dedent();
+        uwriteln!(self.gen.init, "else:");
+        self.gen.init.indent();
+        uwriteln!(
+            self.gen.init,
+            "module = wasmtime.Module(store.engine, file.read_bytes())"
+        );
+        self.gen.init.dedent();
         uwrite!(
             self.gen.init,
             "instance{} = wasmtime.Instance(store, module, [",

setup.py

@@ -31,6 +31,7 @@
         "Source Code": "https://github.com/bytecodealliance/wasmtime-py",
     },
     packages=['wasmtime'],
+    install_requires=['importlib_resources>=5.10'],
     include_package_data=True,
     package_data={"wasmtime": ["py.typed"]},
     python_requires='>=3.8',

wasmtime/loader.py

@@ -7,14 +7,19 @@
 `your_wasm_file.wasm` and hook it up into Python's module system.
 """
 
+from typing import NoReturn, Iterator, Mapping, Dict
+import io
+import re
 import sys
+import struct
 from pathlib import Path
 from importlib import import_module
-from importlib.abc import Loader, MetaPathFinder
+from importlib.abc import Loader, MetaPathFinder, ResourceReader
 from importlib.machinery import ModuleSpec
 
 from wasmtime import Module, Linker, Store, WasiConfig
 from wasmtime import Func, Table, Global, Memory
+from wasmtime import wat2wasm, bindgen
 
 
 predefined_modules = []
@@ -28,7 +33,10 @@
 linker.allow_shadowing = True
 
 
-class _WasmtimeLoader(Loader):
+_component_bindings: Dict[Path, Mapping[str, bytes]] = {}
+
+
+class _CoreWasmLoader(Loader):
     def create_module(self, spec):  # type: ignore
         return None  # use default module creation semantics
 
@@ -55,16 +63,107 @@ def exec_module(self, module):  # type: ignore
             module.__dict__[wasm_export.name] = item
 
 
+class _PythonLoader(Loader):
+    def __init__(self, resource_reader: ResourceReader):
+        self.resource_reader = resource_reader
+
+    def create_module(self, spec):  # type: ignore
+        return None  # use default module creation semantics
+
+    def exec_module(self, module):  # type: ignore
+        origin = Path(module.__spec__.origin)
+        for component_path, component_files in _component_bindings.items():
+            try:
+                relative_path = str(origin.relative_to(component_path))
+            except ValueError:
+                continue
+            exec(component_files[relative_path], module.__dict__)
+            break
+
+    def get_resource_reader(self, fullname: str) -> ResourceReader:
+        return self.resource_reader
+
+
+class _BindingsResourceReader(ResourceReader):
+    def __init__(self, origin: Path):
+        self.resources = _component_bindings[origin]
+
+    def contents(self) -> Iterator[str]:
+        return iter(self.resources.keys())
+
+    def is_resource(self, path: str) -> bool:
+        return path in self.resources
+
+    def open_resource(self, resource: str) -> io.BytesIO:
+        if resource not in self.resources:
+            raise FileNotFoundError
+        return io.BytesIO(self.resources[resource])
+
+    def resource_path(self, resource: str) -> NoReturn:
+        raise FileNotFoundError # all of our resources are virtual
+
+
 class _WasmtimeMetaPathFinder(MetaPathFinder):
+    @staticmethod
+    def is_component(path: Path, *, binary: bool = True) -> bool:
+        if binary:
+            with path.open("rb") as f:
+                preamble = f.read(8)
+            if len(preamble) != 8:
+                return False
+            magic, version, layer = struct.unpack("<4sHH", preamble)
+            if magic != b"\x00asm":
+                return False
+            if layer != 1: # 0 for core wasm, 1 for components
+                return False
+            return True
+        else:
+            contents = path.read_text()
+            # Not strictly correct, but should be good enough for most cases where
+            # someone is using a component in the textual format.
+            return re.search(r"\s*\(\s*component", contents) is not None
+
+    @staticmethod
+    def load_component(path: Path, *, binary: bool = True) -> Mapping[str, bytes]:
+        component = path.read_bytes()
+        if not binary:
+            component = wat2wasm(component)
+        return bindgen.generate("root", component)
+
     def find_spec(self, fullname, path, target=None):  # type: ignore
         modname = fullname.split(".")[-1]
         if path is None:
             path = sys.path
         for entry in map(Path, path):
+            # Is the requested spec a Python module from generated bindings?
+            if entry in _component_bindings:
+                # Create a spec with a virtual origin pointing into generated bindings.
+                origin = entry / (modname + ".py")
+                return ModuleSpec(fullname, _PythonLoader(_BindingsResourceReader(entry)),
+                                  origin=origin)
+            # Is the requested spec a core Wasm module or a Wasm component?
             for suffix in (".wasm", ".wat"):
+                is_binary = (suffix == ".wasm")
                 origin = entry / (modname + suffix)
                 if origin.exists():
-                    return ModuleSpec(fullname, _WasmtimeLoader(), origin=origin)
+                    # Since the origin is on the filesystem, ensure it has an absolute path.
+                    origin = origin.resolve()
+                    if self.is_component(origin, binary=is_binary):
+                        # Generate bindings for the component and remember them for later.
+                        _component_bindings[origin] = self.load_component(origin, binary=is_binary)
+                        # Create a spec with a virtual origin pointing into generated bindings,
+                        # specifically the `__init__.py` file with the code for the package itself.
+                        spec = ModuleSpec(fullname, _PythonLoader(_BindingsResourceReader(origin)),
+                                          origin=origin / '__init__.py', is_package=True)
+                        # Set the search path to the origin. Importlib will provide both the origin
+                        # and the search locations back to this function as-is, even regardless of
+                        # types, but try to follow existing Python conventions. The `origin` will
+                        # be a key in `_component_bindings`.
+                        spec.submodule_search_locations = [origin]
+                        return spec
+                    else:
+                        # Create a spec with a filesystem origin pointing to thg core Wasm module.
+                        return ModuleSpec(fullname, _CoreWasmLoader(), origin=origin)
         return None