make adding a single module simpler for interpreter users (#4792)

youknowone · web-flow · commit 3737f2a0918d · 2024-03-21T21:48:29.000+09:00
diff --git a/src/interpreter.rs b/src/interpreter.rs
@@ -1,4 +1,4 @@
-use rustpython_vm::{Interpreter, Settings, VirtualMachine};
+use rustpython_vm::{builtins::PyModule, Interpreter, PyRef, Settings, VirtualMachine};
 
 pub type InitHook = Box<dyn FnOnce(&mut VirtualMachine)>;
 
@@ -63,6 +63,15 @@ impl InterpreterConfig {
         self.init_hooks.push(hook);
         self
     }
+    pub fn add_native_module(
+        self,
+        name: String,
+        make_module: fn(&VirtualMachine) -> PyRef<PyModule>,
+    ) -> Self {
+        self.init_hook(Box::new(move |vm| {
+            vm.add_native_module(name, Box::new(make_module))
+        }))
+    }
     #[cfg(feature = "stdlib")]
     pub fn init_stdlib(self) -> Self {
         self.init_hook(Box::new(init_stdlib))
