Coverage for wasmtime/loader.py: 86%

1""" 

2This module is a custom loader for Python which enables importing wasm files 

3directly into Python programs simply through usage of the `import` statement. 

4 

5You can import this module with `import wasmtime.loader` and then afterwards you 

6can `import your_wasm_file` which will automatically compile and instantiate 

7`your_wasm_file.wasm` and hook it up into Python's module system. 

8""" 

9 

10from typing import NoReturn, Iterator, Mapping, Dict 

11import io 

12import re 

13import sys 

14import struct 

15from pathlib import Path 

16from importlib import import_module 

17from importlib.abc import Loader, MetaPathFinder, ResourceReader 

18from importlib.machinery import ModuleSpec 

19 

20from wasmtime import Module, Linker, Store, WasiConfig 

21from wasmtime import Func, Table, Global, Memory 

22from wasmtime import wat2wasm, bindgen 

23 

24 

25predefined_modules = [] 

26store = Store() 

27linker = Linker(store.engine) 

28# TODO: how to configure wasi? 

29store.set_wasi(WasiConfig()) 

30predefined_modules.append("wasi_snapshot_preview1") 

31predefined_modules.append("wasi_unstable") 

32linker.define_wasi() 

33linker.allow_shadowing = True 

34 

35 

36_component_bindings: Dict[Path, Mapping[str, bytes]] = {} 

37 

38 

39class _CoreWasmLoader(Loader): 

40 def create_module(self, spec): # type: ignore 

41 return None # use default module creation semantics 

42 

43 def exec_module(self, module): # type: ignore 

44 wasm_module = Module.from_file(store.engine, module.__spec__.origin) 

45 

46 for wasm_import in wasm_module.imports: 

47 module_name = wasm_import.module 

48 if module_name in predefined_modules: 

49 break 

50 field_name = wasm_import.name 

51 imported_module = import_module(module_name) 

52 item = imported_module.__dict__[field_name] 

53 if not isinstance(item, (Func, Table, Global, Memory)): 

54 item = Func(store, wasm_import.type, item) 

55 linker.define(store, module_name, field_name, item) 

56 

57 exports = linker.instantiate(store, wasm_module).exports(store) 

58 for index, wasm_export in enumerate(wasm_module.exports): 

59 item = exports.by_index[index] 

60 if isinstance(item, Func): 

61 # Partially apply `item` to `store`. 

62 item = (lambda func: lambda *args: func(store, *args))(item) 

63 module.__dict__[wasm_export.name] = item 

64 

65 

66class _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 

87class _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 

106class _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"\x00asm": 

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 

133 def find_spec(self, fullname, path, target=None): # type: ignore 

134 modname = fullname.split(".")[-1] 

135 if path is None: 

136 path = sys.path 

137 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? 

145 for suffix in (".wasm", ".wat"): 

146 is_binary = (suffix == ".wasm") 

147 origin = entry / (modname + suffix) 

148 if origin.exists(): 

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) 

167 return None 

168 

169 

170sys.meta_path.append(_WasmtimeMetaPathFinder())