77`your_wasm_file.wasm` and hook it up into Python's module system.
88"""
99
10+ from typing import NoReturn , Iterator , Mapping , Dict
11+ import io
12+ import re
1013import sys
14+ import struct
1115from pathlib import Path
1216from importlib import import_module
13- from importlib .abc import Loader , MetaPathFinder
17+ from importlib .abc import Loader , MetaPathFinder , ResourceReader
1418from importlib .machinery import ModuleSpec
1519
1620from wasmtime import Module , Linker , Store , WasiConfig
1721from wasmtime import Func , Table , Global , Memory
22+ from wasmtime import wat2wasm , bindgen
1823
1924
2025predefined_modules = []
2833linker .allow_shadowing = True
2934
3035
31- class _WasmtimeLoader (Loader ):
36+ _component_bindings : Dict [Path , Mapping [str , bytes ]] = {}
37+
38+
39+ class _CoreWasmLoader (Loader ):
3240 def create_module (self , spec ): # type: ignore
3341 return None # use default module creation semantics
3442
@@ -55,16 +63,107 @@ def exec_module(self, module): # type: ignore
5563 module .__dict__ [wasm_export .name ] = item
5664
5765
66+ class _PythonLoader (Loader ):
67+ def __init__ (self , resource_reader : ResourceReader ):
68+ self .resource_reader = resource_reader
69+
70+ def create_module (self , spec ): # type: ignore
71+ return None # use default module creation semantics
72+
73+ def exec_module (self , module ): # type: ignore
74+ origin = Path (module .__spec__ .origin )
75+ for component_path , component_files in _component_bindings .items ():
76+ try :
77+ relative_path = str (origin .relative_to (component_path ))
78+ except ValueError :
79+ continue
80+ exec (component_files [relative_path ], module .__dict__ )
81+ break
82+
83+ def get_resource_reader (self , fullname : str ) -> ResourceReader :
84+ return self .resource_reader
85+
86+
87+ class _BindingsResourceReader (ResourceReader ):
88+ def __init__ (self , origin : Path ):
89+ self .resources = _component_bindings [origin ]
90+
91+ def contents (self ) -> Iterator [str ]:
92+ return iter (self .resources .keys ())
93+
94+ def is_resource (self , path : str ) -> bool :
95+ return path in self .resources
96+
97+ def open_resource (self , resource : str ) -> io .BytesIO :
98+ if resource not in self .resources :
99+ raise FileNotFoundError
100+ return io .BytesIO (self .resources [resource ])
101+
102+ def resource_path (self , resource : str ) -> NoReturn :
103+ raise FileNotFoundError # all of our resources are virtual
104+
105+
58106class _WasmtimeMetaPathFinder (MetaPathFinder ):
107+ @staticmethod
108+ def is_component (path : Path , * , binary : bool = True ) -> bool :
109+ if binary :
110+ with path .open ("rb" ) as f :
111+ preamble = f .read (8 )
112+ if len (preamble ) != 8 :
113+ return False
114+ magic , version , layer = struct .unpack ("<4sHH" , preamble )
115+ if magic != b"\x00 asm" :
116+ return False
117+ if layer != 1 : # 0 for core wasm, 1 for components
118+ return False
119+ return True
120+ else :
121+ contents = path .read_text ()
122+ # Not strictly correct, but should be good enough for most cases where
123+ # someone is using a component in the textual format.
124+ return re .search (r"\s*\(\s*component" , contents ) is not None
125+
126+ @staticmethod
127+ def load_component (path : Path , * , binary : bool = True ) -> Mapping [str , bytes ]:
128+ component = path .read_bytes ()
129+ if not binary :
130+ component = wat2wasm (component )
131+ return bindgen .generate ("root" , component )
132+
59133 def find_spec (self , fullname , path , target = None ): # type: ignore
60134 modname = fullname .split ("." )[- 1 ]
61135 if path is None :
62136 path = sys .path
63137 for entry in map (Path , path ):
138+ # Is the requested spec a Python module from generated bindings?
139+ if entry in _component_bindings :
140+ # Create a spec with a virtual origin pointing into generated bindings.
141+ origin = entry / (modname + ".py" )
142+ return ModuleSpec (fullname , _PythonLoader (_BindingsResourceReader (entry )),
143+ origin = origin )
144+ # Is the requested spec a core Wasm module or a Wasm component?
64145 for suffix in (".wasm" , ".wat" ):
146+ is_binary = (suffix == ".wasm" )
65147 origin = entry / (modname + suffix )
66148 if origin .exists ():
67- return ModuleSpec (fullname , _WasmtimeLoader (), origin = origin )
149+ # Since the origin is on the filesystem, ensure it has an absolute path.
150+ origin = origin .resolve ()
151+ if self .is_component (origin , binary = is_binary ):
152+ # Generate bindings for the component and remember them for later.
153+ _component_bindings [origin ] = self .load_component (origin , binary = is_binary )
154+ # Create a spec with a virtual origin pointing into generated bindings,
155+ # specifically the `__init__.py` file with the code for the package itself.
156+ spec = ModuleSpec (fullname , _PythonLoader (_BindingsResourceReader (origin )),
157+ origin = origin / '__init__.py' , is_package = True )
158+ # Set the search path to the origin. Importlib will provide both the origin
159+ # and the search locations back to this function as-is, even regardless of
160+ # types, but try to follow existing Python conventions. The `origin` will
161+ # be a key in `_component_bindings`.
162+ spec .submodule_search_locations = [origin ]
163+ return spec
164+ else :
165+ # Create a spec with a filesystem origin pointing to thg core Wasm module.
166+ return ModuleSpec (fullname , _CoreWasmLoader (), origin = origin )
68167 return None
69168
70169
0 commit comments