diff --git a/crates/vm/src/builtins/dict.rs b/crates/vm/src/builtins/dict.rs index aff7432d06..c2070c2a70 100644 --- a/crates/vm/src/builtins/dict.rs +++ b/crates/vm/src/builtins/dict.rs @@ -1,6 +1,6 @@ use super::{ IterStatus, PositionIterInternal, PyBaseExceptionRef, PyGenericAlias, PyMappingProxy, PySet, - PyStr, PyStrRef, PyTupleRef, PyType, PyTypeRef, set::PySetInner, + PyStrRef, PyTupleRef, PyType, PyTypeRef, set::PySetInner, }; use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult, @@ -568,14 +568,37 @@ impl Py { } } - /// Take a python dictionary and convert it to attributes. - pub fn to_attributes(&self, vm: &VirtualMachine) -> PyAttributes { + pub fn to_attributes_with_nonstring( + &self, + vm: &VirtualMachine, + ) -> PyResult<(PyAttributes, Option)> { let mut attrs = PyAttributes::default(); + let mut non_string: Option = None; for (key, value) in self { - let key: PyRefExact = key.downcast_exact(vm).expect("dict has non-string keys"); - attrs.insert(vm.ctx.intern_str(key), value); + if let Some(key) = key.as_interned_str(vm) { + attrs.insert(key, value); + } else { + let dict = non_string.get_or_insert_with(|| vm.ctx.new_dict()); + dict.set_item(key.as_object(), value, vm)?; + } + } + Ok((attrs, non_string)) + } + + /// Take a python dictionary and convert it to attributes. + pub fn to_attributes(&self, vm: &VirtualMachine) -> PyResult { + let (attrs, non_string) = self.to_attributes_with_nonstring(vm)?; + if let Some(non_string) = non_string { + let (key, _) = non_string + .into_iter() + .next() + .expect("internal error: non_string dict should not be empty when present"); + return Err(vm.new_type_error(format!( + "attributes must be strings, not '{}'", + key.class().name() + ))); } - attrs + Ok(attrs) } pub fn get_item_opt( diff --git a/crates/vm/src/builtins/mappingproxy.rs b/crates/vm/src/builtins/mappingproxy.rs index fb8ff5de9c..b9b17295c8 100644 --- a/crates/vm/src/builtins/mappingproxy.rs +++ b/crates/vm/src/builtins/mappingproxy.rs @@ -1,4 +1,4 @@ -use super::{PyDict, PyDictRef, PyGenericAlias, PyList, PyTuple, PyType, PyTypeRef}; +use super::{PyDictRef, PyGenericAlias, PyList, PyTuple, PyType, PyTypeRef}; use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, atomic_func, @@ -91,13 +91,25 @@ impl Constructor for PyMappingProxy { impl PyMappingProxy { fn get_inner(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult> { match &self.mapping { - MappingProxyInner::Class(class) => Ok(key - .as_interned_str(vm) - .and_then(|key| class.attributes.read().get(key).cloned())), + MappingProxyInner::Class(class) => Self::class_get_inner(class, &key, vm), MappingProxyInner::Mapping(mapping) => mapping.mapping().subscript(&*key, vm).map(Some), } } + fn class_get_inner( + class: &PyTypeRef, + key: &PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult> { + if let Some(key) = key.as_interned_str(vm) { + return Ok(class.attributes.read().get(key).cloned()); + } + if let Some(extra) = class.nonstring_attributes.read().as_ref() { + return extra.get_item_opt(key.as_object(), vm); + } + Ok(None) + } + #[pymethod] fn get( &self, @@ -121,9 +133,15 @@ impl PyMappingProxy { fn _contains(&self, key: &PyObject, vm: &VirtualMachine) -> PyResult { match &self.mapping { - MappingProxyInner::Class(class) => Ok(key - .as_interned_str(vm) - .is_some_and(|key| class.attributes.read().contains_key(key))), + MappingProxyInner::Class(class) => { + if let Some(key) = key.as_interned_str(vm) { + Ok(class.attributes.read().contains_key(key)) + } else if let Some(extra) = class.nonstring_attributes.read().as_ref() { + extra.get_item_opt(key, vm).map(|opt| opt.is_some()) + } else { + Ok(false) + } + } MappingProxyInner::Mapping(mapping) => mapping.to_sequence().contains(key, vm), } } @@ -136,9 +154,7 @@ impl PyMappingProxy { fn to_object(&self, vm: &VirtualMachine) -> PyResult { Ok(match &self.mapping { MappingProxyInner::Mapping(d) => d.as_ref().to_owned(), - MappingProxyInner::Class(c) => { - PyDict::from_attributes(c.attributes.read().clone(), vm)?.to_pyobject(vm) - } + MappingProxyInner::Class(c) => c.as_object_dict(vm)?.to_pyobject(vm), }) } @@ -164,9 +180,7 @@ impl PyMappingProxy { pub fn copy(&self, vm: &VirtualMachine) -> PyResult { match &self.mapping { MappingProxyInner::Mapping(d) => vm.call_method(d, identifier!(vm, copy).as_str(), ()), - MappingProxyInner::Class(c) => { - Ok(PyDict::from_attributes(c.attributes.read().clone(), vm)?.to_pyobject(vm)) - } + MappingProxyInner::Class(c) => Ok(c.as_object_dict(vm)?.to_pyobject(vm)), } } diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs index 68de17f60b..476938d758 100644 --- a/crates/vm/src/builtins/type.rs +++ b/crates/vm/src/builtins/type.rs @@ -1,5 +1,5 @@ use super::{ - PyClassMethod, PyDictRef, PyList, PyStr, PyStrInterned, PyStrRef, PyTupleRef, PyWeak, + PyClassMethod, PyDict, PyDictRef, PyList, PyStr, PyStrInterned, PyStrRef, PyTupleRef, PyWeak, mappingproxy::PyMappingProxy, object, union_, }; use crate::{ @@ -41,6 +41,7 @@ pub struct PyType { pub mro: PyRwLock>, pub subclasses: PyRwLock>>, pub attributes: PyRwLock, + pub nonstring_attributes: PyRwLock>, pub slots: PyTypeSlots, pub heaptype_ext: Option>>, } @@ -56,6 +57,9 @@ unsafe impl crate::object::Traverse for PyType { .iter() .map(|(_, v)| v.traverse(tracer_fn)) .count(); + if let Some(dict) = self.nonstring_attributes.read_recursive().as_ref() { + dict.traverse(tracer_fn); + } } } @@ -347,6 +351,7 @@ impl PyType { mro: PyRwLock::new(mro), subclasses: PyRwLock::default(), attributes: PyRwLock::new(attrs), + nonstring_attributes: PyRwLock::default(), slots, heaptype_ext: Some(Pin::new(Box::new(heaptype_ext))), }, @@ -397,6 +402,7 @@ impl PyType { mro: PyRwLock::new(mro), subclasses: PyRwLock::default(), attributes: PyRwLock::new(attrs), + nonstring_attributes: PyRwLock::default(), slots, heaptype_ext: None, }, @@ -544,6 +550,16 @@ impl PyType { attributes } + pub(crate) fn as_object_dict(&self, vm: &VirtualMachine) -> PyResult { + let dict = PyDict::from_attributes(self.attributes.read().clone(), vm)?.into_ref(&vm.ctx); + if let Some(extra) = self.nonstring_attributes.read().clone() { + for (key, value) in extra { + dict.set_item(key.as_object(), value, vm)?; + } + } + Ok(dict) + } + // bound method for every type pub(crate) fn __new__(zelf: PyRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult { let (subtype, args): (PyRef, FuncArgs) = args.bind(vm)?; @@ -1128,7 +1144,7 @@ impl Constructor for PyType { // If __qualname__ is not provided, we can use the name as default name.clone() }); - let mut attributes = dict.to_attributes(vm); + let (mut attributes, nonstring_attributes) = dict.to_attributes_with_nonstring(vm)?; if let Some(f) = attributes.get_mut(identifier!(vm, __init_subclass__)) && f.class().is(vm.ctx.types.function_type) @@ -1246,6 +1262,10 @@ impl Constructor for PyType { ) .map_err(|e| vm.new_type_error(e))?; + if let Some(extra) = nonstring_attributes { + *typ.nonstring_attributes.write() = Some(extra); + } + if let Some(ref slots) = heaptype_slots { let mut offset = base_member_count; let class_name = typ.name().to_string(); diff --git a/crates/vm/src/object/core.rs b/crates/vm/src/object/core.rs index 60b623ef3e..23177f4430 100644 --- a/crates/vm/src/object/core.rs +++ b/crates/vm/src/object/core.rs @@ -1258,6 +1258,7 @@ pub(crate) fn init_type_hierarchy() -> (PyTypeRef, PyTypeRef, PyTypeRef) { mro: PyRwLock::default(), subclasses: PyRwLock::default(), attributes: PyRwLock::new(Default::default()), + nonstring_attributes: PyRwLock::default(), slots: PyType::make_slots(), heaptype_ext: None, }; @@ -1267,6 +1268,7 @@ pub(crate) fn init_type_hierarchy() -> (PyTypeRef, PyTypeRef, PyTypeRef) { mro: PyRwLock::default(), subclasses: PyRwLock::default(), attributes: PyRwLock::new(Default::default()), + nonstring_attributes: PyRwLock::default(), slots: object::PyBaseObject::make_slots(), heaptype_ext: None, }; @@ -1324,6 +1326,7 @@ pub(crate) fn init_type_hierarchy() -> (PyTypeRef, PyTypeRef, PyTypeRef) { mro: PyRwLock::new(vec![object_type.clone()]), subclasses: PyRwLock::default(), attributes: PyRwLock::default(), + nonstring_attributes: PyRwLock::default(), slots: PyWeak::make_slots(), heaptype_ext: None, }; diff --git a/crates/vm/src/stdlib/thread.rs b/crates/vm/src/stdlib/thread.rs index 3625227939..a01e1b45a1 100644 --- a/crates/vm/src/stdlib/thread.rs +++ b/crates/vm/src/stdlib/thread.rs @@ -298,10 +298,13 @@ pub(crate) mod _thread { kwargs: OptionalArg, vm: &VirtualMachine, ) -> PyResult { + let kw_attrs = kwargs + .map(|k| k.to_attributes(vm)) + .transpose()? + .unwrap_or_default(); let args = FuncArgs::new( args.to_vec(), - kwargs - .map_or_else(Default::default, |k| k.to_attributes(vm)) + kw_attrs .into_iter() .map(|(k, v)| (k.as_str().to_owned(), v)) .collect::(), diff --git a/extra_tests/snippets/builtin_mappingproxy.py b/extra_tests/snippets/builtin_mappingproxy.py index fdd653c408..4fa5c0c26b 100644 --- a/extra_tests/snippets/builtin_mappingproxy.py +++ b/extra_tests/snippets/builtin_mappingproxy.py @@ -23,3 +23,11 @@ def b(): assert A.__dict__.get("not here", "default") == "default" assert A.__dict__.get("a", "default") is A.a assert A.__dict__.get("not here") is None + + +class B(object): + locals()[42] = "abc" + + +assert B.__dict__[42] == "abc" +assert 42 in B.__dict__