Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions crates/derive-impl/src/pymodule.rs
Original file line number Diff line number Diff line change
Expand Up @@ -399,8 +399,8 @@ pub(crate) fn impl_pymodule(args: PyModuleArgs, module_item: Item) -> Result<Tok
pub(crate) fn __init_dict(
vm: &::rustpython_vm::VirtualMachine,
module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
) {
::rustpython_vm::builtins::PyModule::__init_dict_from_def(vm, module);
) -> ::rustpython_vm::PyResult<()> {
::rustpython_vm::builtins::PyModule::__init_dict_from_def(vm, module)
}
},
]);
Expand Down
12 changes: 10 additions & 2 deletions crates/vm/src/builtins/dict.rs
Original file line number Diff line number Diff line change
Expand Up @@ -622,6 +622,15 @@ impl PyDict {
}
}

pub(crate) fn inner_pop_item<K: DictKey + ?Sized>(
&self,
key: &K,
vm: &VirtualMachine,
) -> PyResult<Option<PyObjectRef>> {
let hash = Self::hash_or_unhashable(key, vm)?;
self.entries.remove_if_exists(vm, key, hash)
}

pub fn get_or_insert(
&self,
vm: &VirtualMachine,
Expand Down Expand Up @@ -1340,8 +1349,7 @@ impl Py<PyDict> {
vm: &VirtualMachine,
) -> PyResult<Option<PyObjectRef>> {
if self.exact_dict(vm) {
let hash = PyDict::hash_or_unhashable(key, vm)?;
self.entries.remove_if_exists(vm, key, hash)
self.inner_pop_item(key, vm)
} else {
let value = self.as_object().get_item(key, vm)?;
self.as_object().del_item(key, vm)?;
Expand Down
172 changes: 77 additions & 95 deletions crates/vm/src/builtins/module.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,10 @@ use crate::{
class::PyClassImpl,
convert::ToPyObject,
function::{FuncArgs, PyMethodDef, PySetterValue},
import::{get_spec_file_origin, is_possibly_shadowing_path, is_stdlib_module_name},
import::{
get_spec_file_origin, is_possibly_shadowing_path, is_spec_initializing,
is_stdlib_module_name,
},
types::{GetAttr, Initializer, Representable},
};

Expand Down Expand Up @@ -65,7 +68,7 @@ impl PyModuleDef {
};

// Initialize module dict and methods
PyModule::__init_dict_from_def(vm, &module);
PyModule::__init_dict_from_def(vm, &module)?;
module.__init_methods(vm)?;

Ok(module)
Expand Down Expand Up @@ -121,7 +124,7 @@ impl PyPayload for PyModule {
pub struct ModuleInitArgs {
name: PyStrRef,
#[pyarg(any, default)]
doc: Option<PyStrRef>,
doc: Option<PyObjectRef>,
}

impl PyModule {
Expand Down Expand Up @@ -158,9 +161,9 @@ impl PyModule {
drop(self.native_state.store(Some(state)));
}

pub fn __init_dict_from_def(vm: &VirtualMachine, module: &Py<Self>) {
let doc = module.def.unwrap().doc.map(|doc| doc.to_owned());
module.init_dict(module.name.unwrap(), doc, vm);
pub fn __init_dict_from_def(vm: &VirtualMachine, module: &Py<Self>) -> PyResult<()> {
let doc = module.def.unwrap().doc.map(|doc| doc.as_object());
module.init_dict(module.name.unwrap(), doc, vm)
}
}

Expand All @@ -184,7 +187,10 @@ impl Py<PyModule> {
suppress_cycle: bool,
vm: &VirtualMachine,
) -> PyResult {
let mut attr = self.as_object().generic_getattr_opt(name, None, vm)?;
let mut attr = match self.as_object().generic_getattr_opt(name, None, vm) {
Err(exc) if exc.fast_isinstance(vm.ctx.exceptions.attribute_error) => None,
result => result?,
};
if attr.is_none() {
if let Some(value) =
crate::lazy_import::try_load_submodule(self, name, suppress_cycle, vm)?
Expand Down Expand Up @@ -230,13 +236,13 @@ impl Py<PyModule> {
{
return getattr.call((name.to_owned(),), vm);
}
if suppress_cycle {
return Err(vm.new_attribute_error(format!("module has no attribute '{name}'")));
}
let dict = self.dict();

// Get the raw __name__ object (may be a str subclass)
let mod_name_obj = dict
.get_item_opt(identifier!(vm, __name__), vm)
.ok()
.flatten();
let mod_name_obj = dict.inner_getitem_opt(identifier!(vm, __name__), vm)?;
let mod_name = mod_name_obj
.as_ref()
.and_then(|n| n.downcast_ref::<PyUtf8Str>());
Expand All @@ -246,13 +252,9 @@ impl Py<PyModule> {
return Err(vm.new_attribute_error(format!("module has no attribute '{name}'")));
};

let spec = dict
.get_item_opt(vm.ctx.intern_str("__spec__"), vm)
.ok()
.flatten()
.filter(|s| !vm.is_none(s));
let spec = dict.inner_getitem_opt(vm.ctx.intern_str("__spec__"), vm)?;

let origin = get_spec_file_origin(spec.as_deref(), vm);
let origin = get_spec_file_origin(spec.as_deref(), vm)?;

let is_possibly_shadowing = origin
.as_ref()
Expand All @@ -278,7 +280,7 @@ impl Py<PyModule> {
and prevents importing that standard library module)"
)))
} else {
let is_initializing = PyModule::is_initializing(&dict, vm);
let is_initializing = is_spec_initializing(spec.as_deref(), vm)?;
if is_initializing {
if is_possibly_shadowing {
let origin = origin.as_ref().unwrap();
Expand All @@ -302,7 +304,7 @@ impl Py<PyModule> {
}
} else {
// Check for uninitialized submodule
let submodule_initializing = is_uninitialized_submodule(mod_name, name, vm);
let submodule_initializing = is_uninitialized_submodule(spec.as_deref(), name, vm)?;
if submodule_initializing {
Err(vm.new_attribute_error(format!(
"cannot access submodule '{name}' of module '{mod_display}' \
Expand All @@ -325,21 +327,16 @@ impl Py<PyModule> {
// TODO: should be on PyModule, not Py<PyModule>
pub(crate) fn init_dict(
&self,
name: &'static PyStrInterned,
doc: Option<PyStrRef>,
name: &Py<PyStr>,
doc: Option<&PyObject>,
vm: &VirtualMachine,
) {
) -> PyResult<()> {
let dict = self.dict();
dict.set_item(identifier!(vm, __name__), name.to_object(), vm)
.expect("Failed to set __name__ on module");
dict.set_item(identifier!(vm, __doc__), doc.to_pyobject(vm), vm)
.expect("Failed to set __doc__ on module");
dict.set_item("__package__", vm.ctx.none(), vm)
.expect("Failed to set __package__ on module");
dict.set_item("__loader__", vm.ctx.none(), vm)
.expect("Failed to set __loader__ on module");
dict.set_item("__spec__", vm.ctx.none(), vm)
.expect("Failed to set __spec__ on module");
dict.inner_setitem(identifier!(vm, __name__), name.to_owned().into(), vm)?;
dict.inner_setitem(identifier!(vm, __doc__), doc.to_pyobject(vm), vm)?;
dict.inner_setitem(vm.ctx.intern_str("__package__"), vm.ctx.none(), vm)?;
dict.inner_setitem(vm.ctx.intern_str("__loader__"), vm.ctx.none(), vm)?;
dict.inner_setitem(vm.ctx.intern_str("__spec__"), vm.ctx.none(), vm)
}

pub fn get_attr<'a>(&self, attr_name: impl AsPyStr<'a>, vm: &VirtualMachine) -> PyResult {
Expand All @@ -362,38 +359,42 @@ impl Py<PyModule> {
flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
)]
impl PyModule {
fn load_dict(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyDictRef> {
zelf.as_object()
.get_attr(identifier!(vm, __dict__), vm)?
.downcast::<PyDict>()
.map_err(|_| vm.new_type_error("<module>.__dict__ is not a dictionary"))
}

#[pyslot]
fn slot_new(cls: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult {
Self::new().into_ref_with_type(vm, cls).map(Into::into)
}

#[pymethod]
fn __dir__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<Vec<PyObjectRef>> {
fn __dir__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult {
// First check if __dict__ attribute exists and is actually a dictionary
let dict_attr = zelf.as_object().get_attr(identifier!(vm, __dict__), vm)?;
let dict = dict_attr
.downcast::<PyDict>()
.map_err(|_| vm.new_type_error("<module>.__dict__ is not a dictionary"))?;
let dict = Self::load_dict(zelf, vm)?;
// PEP 562: honor a module-level __dir__ if one is defined
if let Some(dir_func) =
crate::lazy_import::resolved_dict_item(&dict, identifier!(vm, __dir__), vm)?
{
return dir_func.call((), vm)?.try_to_value(vm);
return dir_func.call((), vm);
}
let attrs = dict.into_iter().map(|(k, _v)| k).collect();
Ok(attrs)
Ok(vm.ctx.new_list(attrs).into())
}

#[pygetset]
fn __annotate__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
let dict = zelf.dict();
let dict = Self::load_dict(zelf, vm)?;
// Get __annotate__ from dict; if not present, insert None and return it
// See: module_get_annotate()
if let Some(annotate) = dict.get_item_opt(identifier!(vm, __annotate__), vm)? {
if let Some(annotate) = dict.inner_getitem_opt(identifier!(vm, __annotate__), vm)? {
Ok(annotate)
} else {
let none = vm.ctx.none();
dict.set_item(identifier!(vm, __annotate__), none.clone(), vm)?;
dict.inner_setitem(identifier!(vm, __annotate__), none.clone(), vm)?;
Ok(none)
}
}
Expand All @@ -406,14 +407,14 @@ impl PyModule {
) -> PyResult<()> {
match value {
PySetterValue::Assign(value) => {
let dict = Self::load_dict(zelf, vm)?;
if !vm.is_none(&value) && !value.is_callable() {
return Err(vm.new_type_error("__annotate__ must be callable or None"));
}
let dict = zelf.dict();
dict.set_item(identifier!(vm, __annotate__), value.clone(), vm)?;
dict.inner_setitem(identifier!(vm, __annotate__), value.clone(), vm)?;
// Clear __annotations__ if value is not None
if !vm.is_none(&value) {
dict.del_item(identifier!(vm, __annotations__), vm).ok();
dict.inner_pop_item(identifier!(vm, __annotations__), vm)?;
}
Ok(())
}
Expand All @@ -423,19 +424,23 @@ impl PyModule {

#[pygetset]
fn __annotations__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
let dict = zelf.dict();
let dict = Self::load_dict(zelf, vm)?;

// Check if __annotations__ is already in dict (explicitly set)
if let Some(annotations) = dict.get_item_opt(identifier!(vm, __annotations__), vm)? {
if let Some(annotations) = dict.inner_getitem_opt(identifier!(vm, __annotations__), vm)? {
return Ok(annotations);
}

// Check if module is initializing
let is_initializing = Self::is_initializing(&dict, vm);
let spec = zelf
.dict()
.inner_getitem_opt(vm.ctx.intern_str("__spec__"), vm)?;
let is_initializing = is_spec_initializing(spec.as_deref(), vm)?;
drop(spec);

// PEP 649: Get __annotate__ and call it if callable
let annotations = if let Some(annotate) =
dict.get_item_opt(identifier!(vm, __annotate__), vm)?
dict.inner_getitem_opt(identifier!(vm, __annotate__), vm)?
&& annotate.is_callable()
{
// Call __annotate__(1) where 1 is FORMAT_VALUE
Expand All @@ -455,45 +460,34 @@ impl PyModule {

// Cache result unless module is initializing
if !is_initializing {
dict.set_item(identifier!(vm, __annotations__), annotations.clone(), vm)?;
dict.inner_setitem(identifier!(vm, __annotations__), annotations.clone(), vm)?;
}

Ok(annotations)
}

/// Check if module is initializing via __spec__._initializing
fn is_initializing(dict: &Py<PyDict>, vm: &VirtualMachine) -> bool {
if let Ok(Some(spec)) = dict.get_item_opt(vm.ctx.intern_str("__spec__"), vm)
&& let Ok(initializing) = spec.get_attr(vm.ctx.intern_str("_initializing"), vm)
{
return initializing.try_to_bool(vm).unwrap_or(false);
}
false
}

#[pygetset(setter)]
fn set___annotations__(
zelf: &Py<Self>,
value: PySetterValue,
vm: &VirtualMachine,
) -> PyResult<()> {
let dict = zelf.dict();
match value {
let dict = Self::load_dict(zelf, vm)?;
match &value {
PySetterValue::Assign(value) => {
dict.set_item(identifier!(vm, __annotations__), value, vm)?;
// Clear __annotate__ from dict
dict.del_item(identifier!(vm, __annotate__), vm).ok();
Ok(())
dict.inner_setitem(identifier!(vm, __annotations__), value.clone(), vm)?;
}
PySetterValue::Delete => {
if dict.del_item(identifier!(vm, __annotations__), vm).is_err() {
if dict
.inner_pop_item(identifier!(vm, __annotations__), vm)?
.is_none()
{
return Err(vm.new_attribute_error("__annotations__"));
}
// Also clear __annotate__
dict.del_item(identifier!(vm, __annotate__), vm).ok();
Ok(())
}
}
dict.inner_pop_item(identifier!(vm, __annotate__), vm)?;
Ok(())
}
}

Expand All @@ -507,8 +501,8 @@ impl Initializer for PyModule {
.flags
.has_feature(crate::types::PyTypeFlags::HAS_DICT)
);
zelf.init_dict(vm.ctx.intern_str(args.name.as_wtf8()), args.doc, vm);
Ok(())
let ModuleInitArgs { name, doc } = args;
zelf.init_dict(&name, doc.as_deref(), vm)
}
}

Expand Down Expand Up @@ -538,31 +532,19 @@ pub(crate) fn init(context: &'static Context) {
PyModule::extend_class(context, context.types.module_type);
}

/// Check if {module_name}.{name} is an uninitialized submodule in sys.modules.
/// Check the parent module's spec for an uninitialized submodule.
fn is_uninitialized_submodule(
module_name: Option<&Py<PyUtf8Str>>,
spec: Option<&PyObject>,
name: &Py<PyStr>,
vm: &VirtualMachine,
) -> bool {
let Some(mod_name) = module_name else {
return false;
};

let Ok(sys_modules) = vm.sys_module.get_attr("modules", vm) else {
return false;
) -> PyResult<bool> {
let Some(spec) = spec else {
return Ok(false);
};

let full_name = vm.ctx.new_utf8_str(format!("{}.{name}", mod_name.as_str()));
let Ok(sub_mod) = sys_modules.get_item(&*full_name, vm) else {
return false;
};

let spec = match sub_mod.get_attr("__spec__", vm) {
Ok(s) if !vm.is_none(&s) => s,
_ => return false,
};

spec.get_attr("_initializing", vm)
.and_then(|v| v.try_to_bool(vm))
.unwrap_or(false)
match vm.get_attribute_opt(spec, "_uninitialized_submodules")? {
Some(submodules) => submodules
.sequence_unchecked()
.contains(name.as_object(), vm),
None => Ok(false),
}
}
Loading
Loading