Merge pull request RustPython#4898 from youknowone/pymodule-expose

youknowone · web-flow · commit ce13fc5f0327 · 2023-04-25T04:00:52.000+09:00
module objects' type as PyModule
diff --git a/compiler/ast/asdl_rs.py b/compiler/ast/asdl_rs.py
@@ -474,7 +474,7 @@ class ExtendModuleVisitor(EmitVisitor):
     def visitModule(self, mod):
         depth = 0
         self.emit(
-            "pub fn extend_module_nodes(vm: &VirtualMachine, module: &PyObject) {",
+            "pub fn extend_module_nodes(vm: &VirtualMachine, module: &Py<PyModule>) {",
             depth,
         )
         self.emit("extend_module!(vm, module, {", depth + 1)
diff --git a/derive-impl/src/pymodule.rs b/derive-impl/src/pymodule.rs
@@ -114,7 +114,7 @@ pub fn impl_pymodule(attr: AttributeArgs, module_item: Item) -> Result<TokenStre
         parse_quote! {
             pub(crate) fn extend_module(
                 vm: &::rustpython_vm::VirtualMachine,
-                module: &::rustpython_vm::PyObject,
+                module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
             ) {
                 #module_extend_items
             }
@@ -123,7 +123,7 @@ pub fn impl_pymodule(attr: AttributeArgs, module_item: Item) -> Result<TokenStre
             #[allow(dead_code)]
             pub(crate) fn make_module(
                 vm: &::rustpython_vm::VirtualMachine
-            ) -> ::rustpython_vm::PyObjectRef {
+            ) -> ::rustpython_vm::PyRef<::rustpython_vm::builtins::PyModule> {
                 let module = vm.new_module(MODULE_NAME, vm.ctx.new_dict(), DOC);
                 extend_module(vm, &module);
                 module
diff --git a/src/lib.rs b/src/lib.rs
@@ -99,15 +99,12 @@ fn setup_main_module(vm: &VirtualMachine) -> PyResult<Scope> {
     let main_module = vm.new_module("__main__", scope.globals.clone(), None);
     main_module
         .dict()
-        .and_then(|d| {
-            d.set_item("__annotations__", vm.ctx.new_dict().into(), vm)
-                .ok()
-        })
+        .set_item("__annotations__", vm.ctx.new_dict().into(), vm)
         .expect("Failed to initialize __main__.__annotations__");
 
     vm.sys_module
         .get_attr("modules", vm)?
-        .set_item("__main__", main_module, vm)?;
+        .set_item("__main__", main_module.into(), vm)?;
 
     Ok(scope)
 }
diff --git a/stdlib/src/array.rs b/stdlib/src/array.rs
@@ -1,8 +1,8 @@
 // spell-checker:ignore typecode tofile tolist fromfile
 
-use rustpython_vm::{PyObjectRef, VirtualMachine};
+use rustpython_vm::{builtins::PyModule, PyRef, VirtualMachine};
 
-pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let module = array::make_module(vm);
 
     let array = module
diff --git a/stdlib/src/pyexpat.rs b/stdlib/src/pyexpat.rs
@@ -3,12 +3,12 @@
 *
 */
 
-use crate::vm::{extend_module, PyObjectRef, VirtualMachine};
+use crate::vm::{builtins::PyModule, extend_module, PyRef, VirtualMachine};
 
-pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let module = _pyexpat::make_module(vm);
 
-    extend_module!(vm, module, {
+    extend_module!(vm, &module, {
          "errors" => _errors::make_module(vm),
          "model" => _model::make_module(vm),
     });
diff --git a/stdlib/src/select.rs b/stdlib/src/select.rs
@@ -1,9 +1,10 @@
 use crate::vm::{
-    builtins::PyListRef, PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine,
+    builtins::PyListRef, builtins::PyModule, PyObject, PyObjectRef, PyRef, PyResult, TryFromObject,
+    VirtualMachine,
 };
 use std::{io, mem};
 
-pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     #[cfg(windows)]
     crate::vm::stdlib::nt::init_winsock();
 
diff --git a/stdlib/src/socket.rs b/stdlib/src/socket.rs
@@ -1,8 +1,8 @@
-use crate::vm::{PyObjectRef, VirtualMachine};
+use crate::vm::{builtins::PyModule, PyRef, VirtualMachine};
 #[cfg(feature = "ssl")]
 pub(super) use _socket::{sock_select, timeout_error_msg, PySocket, SelectKind};
 
-pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     #[cfg(windows)]
     crate::vm::stdlib::nt::init_winsock();
     _socket::make_module(vm)
diff --git a/stdlib/src/sqlite.rs b/stdlib/src/sqlite.rs
@@ -8,13 +8,13 @@
 // spell-checker:ignore cantlock commithook foreignkey notnull primarykey gettemppath autoindex convpath
 // spell-checker:ignore dbmoved vnode nbytes
 
-use rustpython_vm::{PyObjectRef, VirtualMachine};
+use rustpython_vm::{builtins::PyModule, AsObject, PyRef, VirtualMachine};
 
 // pub(crate) use _sqlite::make_module;
-pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     // TODO: sqlite version check
     let module = _sqlite::make_module(vm);
-    _sqlite::setup_module(&module, vm);
+    _sqlite::setup_module(module.as_object(), vm);
     module
 }
 
diff --git a/stdlib/src/ssl.rs b/stdlib/src/ssl.rs
@@ -1,6 +1,6 @@
-use crate::vm::{PyObjectRef, VirtualMachine};
+use crate::vm::{builtins::PyModule, PyRef, VirtualMachine};
 
-pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     // if openssl is vendored, it doesn't know the locations of system certificates
     #[cfg(feature = "ssl-vendor")]
     if let None | Some("0") = option_env!("OPENSSL_NO_VENDOR") {
diff --git a/stdlib/src/unicodedata.rs b/stdlib/src/unicodedata.rs
@@ -5,11 +5,11 @@
 // spell-checker:ignore nfkc unistr unidata
 
 use crate::vm::{
-    builtins::PyStr, convert::TryFromBorrowedObject, PyObject, PyObjectRef, PyPayload, PyResult,
-    VirtualMachine,
+    builtins::PyModule, builtins::PyStr, convert::TryFromBorrowedObject, PyObject, PyObjectRef,
+    PyPayload, PyRef, PyResult, VirtualMachine,
 };
 
-pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let module = unicodedata::make_module(vm);
 
     let ucd: PyObjectRef = unicodedata::Ucd::new(unic_ucd_age::UNICODE_VERSION)
diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -85,6 +85,7 @@ pub fn make_frozen(vm: &VirtualMachine, name: &str) -> PyResult<PyRef<PyCode>> {
 pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult {
     let frozen = make_frozen(vm, module_name)?;
     let module = import_codeobj(vm, module_name, frozen, false)?;
+    debug_assert!(module.get_attr(identifier!(vm, __name__), vm).is_ok());
     // TODO: give a correct origname here
     module.set_attr("__origname__", vm.ctx.new_str(module_name.to_owned()), vm)?;
     Ok(module)
@@ -99,8 +100,8 @@ pub fn import_builtin(vm: &VirtualMachine, module_name: &str) -> PyResult {
     })?;
     let module = make_module_func(vm);
     let sys_modules = vm.sys_module.get_attr("modules", vm)?;
-    sys_modules.set_item(module_name, module.clone(), vm)?;
-    Ok(module)
+    sys_modules.set_item(module_name, module.as_object().to_owned(), vm)?;
+    Ok(module.into())
 }
 
 #[cfg(feature = "rustpython-compiler")]
@@ -144,12 +145,12 @@ pub fn import_codeobj(
 
     // Store module in cache to prevent infinite loop with mutual importing libs:
     let sys_modules = vm.sys_module.get_attr("modules", vm)?;
-    sys_modules.set_item(module_name, module.clone(), vm)?;
+    sys_modules.set_item(module_name, module.clone().into(), vm)?;
 
     // Execute main code in module:
     let scope = Scope::with_builtins(None, attrs, vm);
     vm.run_code_obj(code_obj, scope)?;
-    Ok(module)
+    Ok(module.into())
 }
 
 fn remove_importlib_frames_inner(
diff --git a/vm/src/macros.rs b/vm/src/macros.rs
@@ -1,10 +1,8 @@
 #[macro_export]
 macro_rules! extend_module {
     ( $vm:expr, $module:expr, { $($name:expr => $value:expr),* $(,)? }) => {{
-        #[allow(unused_variables)]
-        let module: &$crate::PyObject = &$module;
         $(
-            $vm.__module_set_attr(&module, $name, $value).unwrap();
+            $vm.__module_set_attr($module, $name, $value).unwrap();
         )*
     }};
 }
@@ -49,8 +47,9 @@ macro_rules! py_namespace {
     ( $vm:expr, { $($name:expr => $value:expr),* $(,)* }) => {
         {
             let namespace = $crate::builtins::PyNamespace::new_ref(&$vm.ctx);
+            let obj = $crate::object::AsObject::as_object(&namespace);
             $(
-                $vm.__module_set_attr($crate::object::AsObject::as_object(&namespace), $name, $value).unwrap();
+                obj.generic_setattr($vm.ctx.intern_str($name), $crate::function::PySetterValue::Assign($value.into()), $vm).unwrap();
             )*
             namespace
         }
diff --git a/vm/src/stdlib/ast.rs b/vm/src/stdlib/ast.rs
@@ -6,11 +6,11 @@
 mod gen;
 
 use crate::{
-    builtins::{self, PyStrRef, PyType},
+    builtins::{self, PyModule, PyStrRef, PyType},
     class::{PyClassImpl, StaticType},
     compiler::CompileError,
     convert::ToPyException,
-    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyResult, TryFromObject,
+    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
     VirtualMachine,
 };
 use num_complex::Complex64;
@@ -331,7 +331,7 @@ pub(crate) use _ast::AstNode;
 // Used by builtins::compile()
 pub const PY_COMPILE_FLAG_AST_ONLY: i32 = 0x0400;
 
-pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let module = _ast::make_module(vm);
     gen::extend_module_nodes(vm, &module);
     module
diff --git a/vm/src/stdlib/ast/gen.rs b/vm/src/stdlib/ast/gen.rs
diff --git a/vm/src/stdlib/builtins.rs b/vm/src/stdlib/builtins.rs
@@ -1,7 +1,8 @@
 //! Builtin function definitions.
 //!
 //! Implements the list of [builtin Python functions](https://docs.python.org/3/library/builtins.html).
-use crate::{class::PyClassImpl, PyObjectRef, VirtualMachine};
+use crate::{builtins::PyModule, class::PyClassImpl, Py, VirtualMachine};
+pub use builtins::{ascii, print};
 
 #[pymodule]
 mod builtins {
@@ -933,14 +934,12 @@ mod builtins {
     }
 }
 
-pub use builtins::{ascii, print};
-
-pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) {
+pub fn init_module(vm: &VirtualMachine, module: &Py<PyModule>) {
     let ctx = &vm.ctx;
 
     crate::protocol::VecBuffer::make_class(&vm.ctx);
 
-    builtins::extend_module(vm, &module);
+    builtins::extend_module(vm, module);
 
     let debug_mode: bool = vm.state.settings.optimize == 0;
     extend_module!(vm, module, {
diff --git a/vm/src/stdlib/errno.rs b/vm/src/stdlib/errno.rs
@@ -1,13 +1,12 @@
-use crate::PyObjectRef;
-use crate::VirtualMachine;
+use crate::{builtins::PyModule, PyRef, VirtualMachine};
 
 #[pymodule]
 mod errno {}
 
-pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let module = errno::make_module(vm);
     let errorcode = vm.ctx.new_dict();
-    extend_module!(vm, module, {
+    extend_module!(vm, &module, {
         "errorcode" => errorcode.clone(),
     });
     for (name, code) in ERROR_CODES {
diff --git a/vm/src/stdlib/imp.rs b/vm/src/stdlib/imp.rs
@@ -1,7 +1,10 @@
 use crate::bytecode::frozen_lib::FrozenModule;
-use crate::{builtins::PyBaseExceptionRef, PyObjectRef, VirtualMachine};
+use crate::{
+    builtins::{PyBaseExceptionRef, PyModule},
+    PyRef, VirtualMachine,
+};
 
-pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let module = _imp::make_module(vm);
     lock::extend_module(vm, &module);
     module
@@ -120,7 +123,7 @@ mod _imp {
         if let Ok(module) = sys_modules.get_item(&*name, vm) {
             Ok(module)
         } else if let Some(make_module_func) = vm.state.module_inits.get(name.as_str()) {
-            Ok(make_module_func(vm))
+            Ok(make_module_func(vm).into())
         } else {
             Ok(vm.ctx.none())
         }
diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs
@@ -11,8 +11,9 @@ cfg_if::cfg_if! {
 
 use crate::{
     builtins::PyBaseExceptionRef,
+    builtins::PyModule,
     convert::{IntoPyException, ToPyException, ToPyObject},
-    PyObjectRef, PyResult, TryFromObject, VirtualMachine,
+    PyObjectRef, PyRef, PyResult, TryFromObject, VirtualMachine,
 };
 pub use _io::io_open as open;
 
@@ -44,7 +45,7 @@ impl IntoPyException for std::io::Error {
     }
 }
 
-pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let ctx = &vm.ctx;
 
     let module = _io::make_module(vm);
@@ -55,7 +56,7 @@ pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let unsupported_operation = _io::UNSUPPORTED_OPERATION
         .get_or_init(|| _io::make_unsupportedop(ctx))
         .clone();
-    extend_module!(vm, module, {
+    extend_module!(vm, &module, {
         "UnsupportedOperation" => unsupported_operation,
         "BlockingIOError" => ctx.exceptions.blocking_io_error.to_owned(),
     });
diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs
@@ -48,10 +48,10 @@ mod winapi;
 #[cfg(windows)]
 mod winreg;
 
-use crate::{PyObjectRef, VirtualMachine};
+use crate::{builtins::PyModule, PyRef, VirtualMachine};
 use std::{borrow::Cow, collections::HashMap};
 
-pub type StdlibInitFunc = Box<py_dyn_fn!(dyn Fn(&VirtualMachine) -> PyObjectRef)>;
+pub type StdlibInitFunc = Box<py_dyn_fn!(dyn Fn(&VirtualMachine) -> PyRef<PyModule>)>;
 pub type StdlibMap = HashMap<Cow<'static, str>, StdlibInitFunc, ahash::RandomState>;
 
 pub fn get_module_inits() -> StdlibMap {
diff --git a/vm/src/stdlib/nt.rs b/vm/src/stdlib/nt.rs
@@ -1,8 +1,8 @@
-use crate::{PyObjectRef, VirtualMachine};
+use crate::{builtins::PyModule, PyRef, VirtualMachine};
 
 pub use module::raw_set_handle_inheritable;
 
-pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let module = module::make_module(vm);
     super::os::extend_module(vm, &module);
     module
diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs
@@ -1,9 +1,9 @@
 use crate::{
-    builtins::{PyBaseExceptionRef, PySet},
+    builtins::{PyBaseExceptionRef, PyModule, PySet},
     common::crt_fd::Fd,
     convert::IntoPyException,
     function::{ArgumentError, FromArgs, FsPath, FuncArgs},
-    AsObject, PyObject, PyObjectRef, PyPayload, PyResult, TryFromObject, VirtualMachine,
+    AsObject, Py, PyObjectRef, PyPayload, PyResult, TryFromObject, VirtualMachine,
 };
 use std::{
     ffi, fs, io,
@@ -1709,7 +1709,7 @@ impl SupportFunc {
     }
 }
 
-pub fn extend_module(vm: &VirtualMachine, module: &PyObject) {
+pub fn extend_module(vm: &VirtualMachine, module: &Py<PyModule>) {
     _os::extend_module(vm, module);
 
     let support_funcs = _os::support_funcs();
diff --git a/vm/src/stdlib/posix.rs b/vm/src/stdlib/posix.rs
@@ -1,4 +1,4 @@
-use crate::{PyObjectRef, VirtualMachine};
+use crate::{builtins::PyModule, PyRef, VirtualMachine};
 use std::os::unix::io::RawFd;
 
 pub fn raw_set_inheritable(fd: RawFd, inheritable: bool) -> nix::Result<()> {
@@ -12,7 +12,7 @@ pub fn raw_set_inheritable(fd: RawFd, inheritable: bool) -> nix::Result<()> {
     Ok(())
 }
 
-pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let module = module::make_module(vm);
     super::os::extend_module(vm, &module);
     module
diff --git a/vm/src/stdlib/posix_compat.rs b/vm/src/stdlib/posix_compat.rs
@@ -1,8 +1,8 @@
 //! `posix` compatible module for `not(any(unix, windows))`
 
-use crate::{PyObjectRef, VirtualMachine};
+use crate::{builtins::PyModule, PyObjectRef, PyRef, VirtualMachine};
 
-pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let module = module::make_module(vm);
     super::os::extend_module(vm, &module);
     module
diff --git a/vm/src/stdlib/signal.rs b/vm/src/stdlib/signal.rs
diff --git a/vm/src/stdlib/sys.rs b/vm/src/stdlib/sys.rs
diff --git a/vm/src/stdlib/time.rs b/vm/src/stdlib/time.rs
diff --git a/vm/src/stdlib/winreg.rs b/vm/src/stdlib/winreg.rs
diff --git a/vm/src/vm/mod.rs b/vm/src/vm/mod.rs
diff --git a/vm/src/vm/vm_new.rs b/vm/src/vm/vm_new.rs
diff --git a/wasm/lib/src/vm_class.rs b/wasm/lib/src/vm_class.rs

Original file line number	Diff line number	Diff line change
`@@ -474,7 +474,7 @@ class ExtendModuleVisitor(EmitVisitor):`
`474`	`474`	`def visitModule(self, mod):`
`475`	`475`	`depth = 0`
`476`	`476`	`self.emit(`
`477`		`- "pub fn extend_module_nodes(vm: &VirtualMachine, module: &PyObject) {",`
	`477`	`+ "pub fn extend_module_nodes(vm: &VirtualMachine, module: &Py<PyModule>) {",`
`478`	`478`	`depth,`
`479`	`479`	`)`
`480`	`480`	`self.emit("extend_module!(vm, module, {", depth + 1)`