diff --git a/Lib/copy.py b/Lib/copy.py index c64fc076179..6149301ad13 100644 --- a/Lib/copy.py +++ b/Lib/copy.py @@ -100,14 +100,14 @@ def copy(x): return _reconstruct(x, None, *rv) -_copy_atomic_types = {types.NoneType, int, float, bool, complex, str, tuple, - bytes, frozenset, type, range, slice, property, +_copy_atomic_types = frozenset({types.NoneType, int, float, bool, complex, str, tuple, + bytes, frozendict, frozenset, type, range, slice, property, types.BuiltinFunctionType, types.EllipsisType, types.NotImplementedType, types.FunctionType, types.CodeType, - weakref.ref, super} -_copy_builtin_containers = {list, dict, set, bytearray} + weakref.ref, super}) +_copy_builtin_containers = frozenset({list, dict, set, bytearray}) -def deepcopy(x, memo=None, _nil=[]): +def deepcopy(x, memo=None): """Deep copy operation on arbitrary Python objects. See the module's __doc__ string for more info. @@ -122,8 +122,8 @@ def deepcopy(x, memo=None, _nil=[]): if memo is None: memo = {} else: - y = memo.get(d, _nil) - if y is not _nil: + y = memo.get(d, None) + if y is not None: return y copier = _deepcopy_dispatch.get(cls) @@ -162,9 +162,9 @@ def deepcopy(x, memo=None, _nil=[]): _keep_alive(x, memo) # Make sure x lives at least as long as d return y -_atomic_types = {types.NoneType, types.EllipsisType, types.NotImplementedType, +_atomic_types = frozenset({types.NoneType, types.EllipsisType, types.NotImplementedType, int, float, bool, complex, bytes, str, types.CodeType, type, range, - types.BuiltinFunctionType, types.FunctionType, weakref.ref, property} + types.BuiltinFunctionType, types.FunctionType, weakref.ref, property}) _deepcopy_dispatch = d = {} @@ -203,6 +203,21 @@ def _deepcopy_dict(x, memo, deepcopy=deepcopy): return y d[dict] = _deepcopy_dict +def _deepcopy_frozendict(x, memo, deepcopy=deepcopy): + y = {} + for key, value in x.items(): + y[deepcopy(key, memo)] = deepcopy(value, memo) + + # We're not going to put the frozendict in the memo, but it's still + # important we check for it, in case the frozendict contains recursive + # mutable structures. + try: + return memo[id(x)] + except KeyError: + pass + return frozendict(y) +d[frozendict] = _deepcopy_frozendict + def _deepcopy_method(x, memo): # Copy instance methods return type(x)(x.__func__, deepcopy(x.__self__, memo)) d[types.MethodType] = _deepcopy_method @@ -230,7 +245,7 @@ def _reconstruct(x, memo, func, args, *, deepcopy=deepcopy): deep = memo is not None if deep and args: - args = (deepcopy(arg, memo) for arg in args) + args = [deepcopy(arg, memo) for arg in args] y = func(*args) if deep: memo[id(x)] = y diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py index 41e4087bacb..75e2afa7ae6 100644 --- a/Lib/test/pickletester.py +++ b/Lib/test/pickletester.py @@ -4394,8 +4394,6 @@ def test_fast_save_enter_frozenset(self): def test_fast_save_enter_set(self): self.fast_save_enter(lambda i: set([i])) - # TODO: RUSTPYTHON; frozendict support: https://github.com/RustPython/RustPython/pull/8951 - @unittest.expectedFailure def test_fast_save_enter_frozendict(self): if self.py_version < (3, 15): self.skipTest('need frozendict') @@ -4443,8 +4441,6 @@ def test_deep_nested_struct_set(self): depth=FAST_NESTING_LIMIT+1, compare_equal=False) - # TODO: RUSTPYTHON; frozendict support: https://github.com/RustPython/RustPython/pull/8951 - @unittest.expectedFailure def test_deep_nested_struct_frozendict(self): if self.py_version < (3, 15): self.skipTest('need frozendict') diff --git a/Lib/test/test_dict.py b/Lib/test/test_dict.py index 195466e6658..8f8ebc15feb 100644 --- a/Lib/test/test_dict.py +++ b/Lib/test/test_dict.py @@ -951,7 +951,6 @@ def test_empty_presized_dict_in_freelist(self): 'f': None, 'g': None, 'h': None} d = {} - @unittest.expectedFailure # TODO: RUSTPYTHON def test_container_iterator(self): # Bug #3680: tp_traverse was not implemented for dictiter and # dictview objects. @@ -1264,7 +1263,6 @@ def test_free_after_iterating(self): support.check_free_after_iterating(self, lambda d: iter(d.values()), dict) support.check_free_after_iterating(self, lambda d: iter(d.items()), dict) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_equal_operator_modifying_operand(self): # test fix for seg fault reported in bpo-27945 part 3. class X(): diff --git a/Lib/test/test_doctest/test_doctest.py b/Lib/test/test_doctest/test_doctest.py index 1945cd3f0dd..b9f67f07ebc 100644 --- a/Lib/test/test_doctest/test_doctest.py +++ b/Lib/test/test_doctest/test_doctest.py @@ -742,7 +742,7 @@ def non_Python_modules(): r""" >>> import builtins >>> tests = doctest.DocTestFinder().find(builtins) - >>> 750 < len(tests) < 800 # approximate number of objects with docstrings + >>> 750 < len(tests) < 850 # approximate number of objects with docstrings True >>> real_tests = [t for t in tests if len(t.examples) > 0] >>> len(real_tests) # objects that actually have doctests diff --git a/Lib/test/test_dynamic.py b/Lib/test/test_dynamic.py index 948a2029f1e..1df010732c1 100644 --- a/Lib/test/test_dynamic.py +++ b/Lib/test/test_dynamic.py @@ -135,7 +135,7 @@ def test_eval_gives_lambda_custom_globals(self): self.assertEqual(foo(), 7) - @unittest.expectedFailure # TODO: RUSTPYTHON + @unittest.expectedFailureIf(sys.getrecursionlimit() == 256, "TODO: RUSTPYTHON; default debug recursion limit is 256") @skip_wasi_stack_overflow() def test_load_global_specialization_failure_keeps_oparg(self): # https://github.com/python/cpython/issues/91625 diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py index 71cd4baa942..1b85eceac17 100644 --- a/Lib/test/test_pickle.py +++ b/Lib/test/test_pickle.py @@ -107,6 +107,10 @@ def loads(self, buf, **kwds): def test_c_methods(self): return super().test_c_methods() + @unittest.expectedFailureIf(sys.getrecursionlimit() == 256, "TODO: RUSTPYTHON; default debug recursion limit is 256") + def test_deep_nested_struct_frozendict(self): + return super().test_deep_nested_struct_frozendict() + @@ -408,6 +412,10 @@ class CPicklerTests(PyPicklerTests): pickler = _pickle.Pickler unpickler = _pickle.Unpickler + @unittest.expectedSuccess # TODO: RUSTPYTHON; native pickler passes at the default debug recursion limit + def test_deep_nested_struct_frozendict(self): + return super().test_deep_nested_struct_frozendict() + class CPersPicklerTests(PyPersPicklerTests): pickler = _pickle.Pickler unpickler = _pickle.Unpickler @@ -421,6 +429,10 @@ class CDumpPickle_LoadPickle(PyPicklerTests): pickler = _pickle.Pickler unpickler = pickle._Unpickler + @unittest.expectedSuccess # TODO: RUSTPYTHON; native pickler passes at the default debug recursion limit + def test_deep_nested_struct_frozendict(self): + return super().test_deep_nested_struct_frozendict() + class DumpPickle_CLoadPickle(PyPicklerTests): pickler = pickle._Pickler unpickler = _pickle.Unpickler diff --git a/Lib/test/test_unittest/testmock/testhelpers.py b/Lib/test/test_unittest/testmock/testhelpers.py index bcae14c5b3c..0e82c723ec3 100644 --- a/Lib/test/test_unittest/testmock/testhelpers.py +++ b/Lib/test/test_unittest/testmock/testhelpers.py @@ -45,7 +45,6 @@ def test_any_and_datetime(self): mock.assert_called_with(ANY, foo=ANY) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_any_mock_calls_comparison_order(self): mock = Mock() class Foo(object): diff --git a/crates/capi/src/ceval.rs b/crates/capi/src/ceval.rs index ace7896163c..f571ea0dbb7 100644 --- a/crates/capi/src/ceval.rs +++ b/crates/capi/src/ceval.rs @@ -3,7 +3,7 @@ use crate::pystate::with_vm; use crate::unicodeobject::decode_fsdefault_and_size; use crate::util::{CStrExt, FfiPtrExt}; use core::ffi::{CStr, c_char, c_int}; -use rustpython_vm::builtins::{PyCode, PyDict}; +use rustpython_vm::builtins::{PyAnyDictRef, PyCode}; use rustpython_vm::function::ArgMapping; use rustpython_vm::scope::Scope; use rustpython_vm::{AsObject, PyObject, TryFromObject}; @@ -32,12 +32,15 @@ pub unsafe extern "C" fn PyEval_EvalCode( ) -> *mut PyObject { with_vm(|vm| { let code = unsafe { co.assume_borrowed_and_cast::(vm) }?; - let globals = unsafe { globals.assume_borrowed_and_cast::(vm) }?; + let globals = + PyAnyDictRef::try_from_object(vm, unsafe { globals.assume_borrowed() }.to_owned())?; let locals = unsafe { locals.assume_borrowed_or_opt() } .map(|obj| ArgMapping::try_from_object(vm, obj.to_owned())) .transpose()?; - let scope = Scope::with_builtins(locals, globals.to_owned(), vm); + let locals = locals.or_else(|| Some(ArgMapping::from_anydict_exact(globals.clone()))); + // Unlike Python's eval/exec, this API does not insert __builtins__. + let scope = Scope::new(locals, globals); vm.run_code_obj(code.to_owned(), scope) }) diff --git a/crates/capi/src/dictobject.rs b/crates/capi/src/dictobject.rs index 2c3eedff3ba..3ca9304e330 100644 --- a/crates/capi/src/dictobject.rs +++ b/crates/capi/src/dictobject.rs @@ -4,9 +4,8 @@ use crate::pystate::with_vm; use crate::util::{CStrExt, FfiPtrExt}; use core::ffi::{c_char, c_int}; use core::ptr::NonNull; -use rustpython_vm::AsObject; -use rustpython_vm::PyPayload; -use rustpython_vm::builtins::PyDict; +use rustpython_vm::builtins::{PyAnyDictRef, PyDict, PyFrozenDict}; +use rustpython_vm::{AsObject, Py, PyPayload, PyResult, VirtualMachine}; define_py_check!(fn PyDict_Check, types.dict_type); define_py_check!(exact fn PyDict_CheckExact, types.dict_type); @@ -14,6 +13,27 @@ define_py_check!(fn PyDictKeys_Check, types.dict_keys_type); define_py_check!(fn PyDictValues_Check, types.dict_values_type); define_py_check!(fn PyDictItems_Check, types.dict_items_type); +fn any_dict(dict: &PyObject, vm: &VirtualMachine) -> PyResult { + PyAnyDictRef::from_object(dict) + .ok_or_else(|| vm.new_system_error("bad argument to internal function")) +} + +fn writable_dict<'a>( + dict: &'a PyObject, + operation: &str, + vm: &VirtualMachine, +) -> PyResult<&'a Py> { + dict.downcast_ref::().ok_or_else(|| { + if dict.downcast_ref::().is_some() { + vm.new_type_error(format!( + "frozendict object does not support item {operation}" + )) + } else { + vm.new_system_error("bad argument to internal function") + } + }) +} + #[unsafe(no_mangle)] pub extern "C" fn PyDict_New() -> *mut PyObject { with_vm(|vm| vm.ctx.new_dict()) @@ -21,10 +41,10 @@ pub extern "C" fn PyDict_New() -> *mut PyObject { #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_Clear(dict: *mut PyObject) { - with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; - dict.clear(); - Ok(()) + with_vm(|_vm| { + if let Some(dict) = unsafe { dict.assume_borrowed() }.downcast_ref::() { + dict.clear(); + } }) } @@ -35,7 +55,7 @@ pub unsafe extern "C" fn PyDict_SetItem( val: *mut PyObject, ) -> c_int { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = writable_dict(unsafe { dict.assume_borrowed() }, "assignment", vm)?; let key = unsafe { key.assume_borrowed() }; let value = unsafe { val.assume_borrowed() }.to_owned(); dict.inner_setitem(key, value, vm) @@ -49,8 +69,8 @@ pub unsafe extern "C" fn PyDict_SetItemString( val: *mut PyObject, ) -> c_int { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; let key = unsafe { key.try_as_str(vm) }?; + let dict = writable_dict(unsafe { dict.assume_borrowed() }, "assignment", vm)?; let value = unsafe { val.assume_borrowed() }.to_owned(); dict.inner_setitem(key, value, vm) }) @@ -59,12 +79,14 @@ pub unsafe extern "C" fn PyDict_SetItemString( #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_GetItem(dict: *mut PyObject, key: *mut PyObject) -> *mut PyObject { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let Some(dict) = PyAnyDictRef::from_object(unsafe { dict.assume_borrowed() }) else { + return core::ptr::null_mut(); + }; let key = unsafe { key.assume_borrowed() }; match dict.inner_getitem_opt(key, vm) { - Ok(Some(value)) => Ok(value.as_object().as_raw().cast_mut()), - Ok(None) | Err(_) => Ok(core::ptr::null_mut()), + Ok(Some(value)) => value.as_object().as_raw().cast_mut(), + Ok(None) | Err(_) => core::ptr::null_mut(), } }) } @@ -75,7 +97,7 @@ pub unsafe extern "C" fn PyDict_GetItemWithError( key: *mut PyObject, ) -> *mut PyObject { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = any_dict(unsafe { dict.assume_borrowed() }, vm)?; let key = unsafe { key.assume_borrowed() }; if let Some(value) = dict.inner_getitem_opt(key, vm)? { @@ -92,12 +114,16 @@ pub unsafe extern "C" fn PyDict_GetItemString( key: *const c_char, ) -> *mut PyObject { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; - let key = unsafe { key.try_as_str(vm) }?; + let Ok(key) = (unsafe { key.try_as_str(vm) }) else { + return core::ptr::null_mut(); + }; + let Some(dict) = PyAnyDictRef::from_object(unsafe { dict.assume_borrowed() }) else { + return core::ptr::null_mut(); + }; - match dict.inner_getitem_opt(key, vm)? { - Some(value) => Ok(value.as_object().as_raw().cast_mut()), - None => Ok(core::ptr::null_mut()), + match dict.inner_getitem_opt(key, vm) { + Ok(Some(value)) => value.as_object().as_raw().cast_mut(), + Ok(None) | Err(_) => core::ptr::null_mut(), } }) } @@ -112,8 +138,8 @@ pub unsafe extern "C" fn PyDict_GetItemStringRef( unsafe { *result = core::ptr::null_mut(); } - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; let key = unsafe { key.try_as_str(vm) }?; + let dict = any_dict(unsafe { dict.assume_borrowed() }, vm)?; if let Some(value) = dict.inner_getitem_opt(key, vm)? { unsafe { @@ -134,7 +160,7 @@ pub unsafe extern "C" fn PyDict_GetItemRef( ) -> c_int { with_vm(|vm| { unsafe { *result = core::ptr::null_mut() }; - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = any_dict(unsafe { dict.assume_borrowed() }, vm)?; let key = unsafe { key.assume_borrowed() }; if let Some(value) = dict.inner_getitem_opt(key, vm)? { @@ -165,7 +191,7 @@ pub unsafe extern "C" fn PyDict_SetDefaultRef( result.write(core::ptr::null_mut()); } } - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = writable_dict(unsafe { dict.assume_borrowed() }, "assignment", vm)?; let key = unsafe { key.assume_borrowed() }; if let Some(value) = dict.inner_getitem_opt(key, vm)? { @@ -191,7 +217,9 @@ pub unsafe extern "C" fn PyDict_SetDefaultRef( #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_Size(dict: *mut PyObject) -> isize { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = unsafe { dict.assume_borrowed_or_opt() } + .ok_or_else(|| vm.new_system_error("bad argument to internal function"))?; + let dict = any_dict(dict, vm)?; Ok(dict.__len__()) }) } @@ -199,7 +227,7 @@ pub unsafe extern "C" fn PyDict_Size(dict: *mut PyObject) -> isize { #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_Contains(dict: *mut PyObject, key: *mut PyObject) -> c_int { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = any_dict(unsafe { dict.assume_borrowed() }, vm)?; let key = unsafe { key.assume_borrowed() }; Ok(dict.inner_getitem_opt(key, vm)?.is_some()) }) @@ -208,7 +236,9 @@ pub unsafe extern "C" fn PyDict_Contains(dict: *mut PyObject, key: *mut PyObject #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_Copy(dict: *mut PyObject) -> *mut PyObject { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = unsafe { dict.assume_borrowed_or_opt() } + .and_then(|dict| dict.downcast_ref::()) + .ok_or_else(|| vm.new_system_error("bad argument to internal function"))?; Ok(dict.copy().into_ref(&vm.ctx)) }) } @@ -216,8 +246,12 @@ pub unsafe extern "C" fn PyDict_Copy(dict: *mut PyObject) -> *mut PyObject { #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_DelItem(dict: *mut PyObject, key: *mut PyObject) -> c_int { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = unsafe { dict.assume_borrowed() }; let key = unsafe { key.assume_borrowed() }; + if dict.downcast_ref::().is_some() { + PyFrozenDict::key_hash(key, vm)?; + } + let dict = writable_dict(dict, "deletion", vm)?; dict.del_item(key, vm) }) } @@ -225,8 +259,8 @@ pub unsafe extern "C" fn PyDict_DelItem(dict: *mut PyObject, key: *mut PyObject) #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_DelItemString(dict: *mut PyObject, key: *const c_char) -> c_int { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; let key = unsafe { key.try_as_str(vm) }?; + let dict = writable_dict(unsafe { dict.assume_borrowed() }, "deletion", vm)?; dict.del_item(key, vm) }) } @@ -234,7 +268,9 @@ pub unsafe extern "C" fn PyDict_DelItemString(dict: *mut PyObject, key: *const c #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_Items(dict: *mut PyObject) -> *mut PyObject { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = unsafe { dict.assume_borrowed_or_opt() } + .ok_or_else(|| vm.new_system_error("bad argument to internal function"))?; + let dict = any_dict(dict, vm)?; let items = dict .items_vec() .into_iter() @@ -247,7 +283,9 @@ pub unsafe extern "C" fn PyDict_Items(dict: *mut PyObject) -> *mut PyObject { #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_Keys(dict: *mut PyObject) -> *mut PyObject { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = unsafe { dict.assume_borrowed_or_opt() } + .ok_or_else(|| vm.new_system_error("bad argument to internal function"))?; + let dict = any_dict(dict, vm)?; Ok(vm.ctx.new_list(dict.keys_vec())) }) } @@ -255,7 +293,9 @@ pub unsafe extern "C" fn PyDict_Keys(dict: *mut PyObject) -> *mut PyObject { #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_Values(dict: *mut PyObject) -> *mut PyObject { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = unsafe { dict.assume_borrowed_or_opt() } + .ok_or_else(|| vm.new_system_error("bad argument to internal function"))?; + let dict = any_dict(dict, vm)?; Ok(vm.ctx.new_list(dict.values_vec())) }) } @@ -267,7 +307,7 @@ pub unsafe extern "C" fn PyDict_Merge( override_: c_int, ) -> c_int { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = writable_dict(unsafe { dict.assume_borrowed() }, "assignment", vm)?; let other = unsafe { other.assume_borrowed() }.to_owned(); if override_ != 0 { dict.merge_object(other, vm) @@ -280,7 +320,7 @@ pub unsafe extern "C" fn PyDict_Merge( #[unsafe(no_mangle)] pub unsafe extern "C" fn PyDict_Update(dict: *mut PyObject, other: *mut PyObject) -> c_int { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = writable_dict(unsafe { dict.assume_borrowed() }, "assignment", vm)?; let other = unsafe { other.assume_borrowed() }.to_owned(); dict.merge_object(other, vm) }) @@ -293,7 +333,7 @@ pub unsafe extern "C" fn PyDict_MergeFromSeq2( override_: c_int, ) -> c_int { with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + let dict = writable_dict(unsafe { dict.assume_borrowed() }, "assignment", vm)?; let seq2 = unsafe { seq2.assume_borrowed() }; dict.merge_from_seq2(seq2, override_ != 0, vm) }) @@ -306,8 +346,10 @@ pub unsafe extern "C" fn PyDict_Next( key: *mut *mut PyObject, value: *mut *mut PyObject, ) -> c_int { - with_vm(|vm| { - let dict = unsafe { dict.assume_borrowed_and_cast::(vm) }?; + with_vm(|_vm| { + let Some(dict) = PyAnyDictRef::from_object(unsafe { dict.assume_borrowed() }) else { + return false; + }; let index = unsafe { *pos } as usize; if let Some((next_pos, k, v)) = dict.next_entry(index) { @@ -320,9 +362,9 @@ pub unsafe extern "C" fn PyDict_Next( value.write(v.as_object().as_raw().cast_mut()); } } - Ok(true) + true } else { - Ok(false) + false } }) } diff --git a/crates/capi/src/pyframe.rs b/crates/capi/src/pyframe.rs index 2214ee55820..c3e81c235b3 100644 --- a/crates/capi/src/pyframe.rs +++ b/crates/capi/src/pyframe.rs @@ -2,9 +2,9 @@ use crate::PyObject; use crate::pystate::with_vm; use crate::util::FfiPtrExt; use core::ffi::c_int; -use rustpython_vm::Py; use rustpython_vm::builtins::PyCode; use rustpython_vm::frame::FrameObject; +use rustpython_vm::{Py, PyObjectRef}; pub type PyFrameObject = Py; pub type PyCodeObject = Py; @@ -26,7 +26,7 @@ pub unsafe extern "C" fn PyFrame_GetLocals(frame: *mut PyFrameObject) -> *mut Py #[unsafe(no_mangle)] pub unsafe extern "C" fn PyFrame_GetGlobals(frame: *mut PyFrameObject) -> *mut PyObject { - with_vm(|_vm| Ok(unsafe { frame.assume_borrowed() }.f_globals())) + with_vm(|_vm| PyObjectRef::from(unsafe { frame.assume_borrowed() }.f_globals())) } #[unsafe(no_mangle)] diff --git a/crates/vm/src/anystr.rs b/crates/vm/src/anystr.rs index b9ce612500b..1b7b20641e1 100644 --- a/crates/vm/src/anystr.rs +++ b/crates/vm/src/anystr.rs @@ -47,6 +47,19 @@ pub(crate) struct StartsEndsWithArgs { } impl StartsEndsWithArgs { + pub(crate) fn parts(&self, len: usize) -> (&PyObject, Option>) { + let range = if self.start.is_some() || self.end.is_some() { + Some(adjust_indices( + self.start.as_ref().map(AsRef::as_ref), + self.end.as_ref().map(AsRef::as_ref), + len, + )) + } else { + None + }; + (&self.affix, range) + } + pub(crate) fn get_value(self, len: usize) -> (PyObjectRef, Option>) { let range = if self.start.is_some() || self.end.is_some() { Some(adjust_indices( diff --git a/crates/vm/src/builtins/bytearray.rs b/crates/vm/src/builtins/bytearray.rs index 5afeec62833..d6d424799e8 100644 --- a/crates/vm/src/builtins/bytearray.rs +++ b/crates/vm/src/builtins/bytearray.rs @@ -6,7 +6,7 @@ use super::{ use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine, - anystr::{self, AnyStr}, + anystr::{self, AnyStr, StringRange}, atomic_func, byte::{bytearray_extend_from_object, bytearray_from_object, value_from_object}, bytes_inner::{ @@ -23,7 +23,7 @@ use crate::{ }, }, convert::{ToPyObject, ToPyResult}, - function::{ArgBytesLike, FuncArgs, PyComparisonValue, PySsize}, + function::{ArgBytesLike, FuncArgs, OptionalArg, PyComparisonValue, PySsize}, protocol::{ BufferDescriptor, BufferFlags, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods, @@ -65,6 +65,21 @@ impl PyPayload for PyByteArray { } } +// Acquisitions can call Python, and releases must follow acquisition order. +// No bytearray lock may be held when this owner is dropped. +#[derive(Default)] +struct ByteArrayBuffers { + first: Option, + second: Option, +} + +impl Drop for ByteArrayBuffers { + fn drop(&mut self) { + drop(self.first.take()); + drop(self.second.take()); + } +} + /// Fill bytearray class methods dictionary. pub(crate) fn init(context: &'static Context) { PyByteArray::extend_class(context, context.types.bytearray_type); @@ -92,6 +107,105 @@ impl PyByteArray { PyDetachingRwLockWriteGuard::map(self.inner.write(), |inner| &mut inner.elements) } + fn with_export(&self, f: impl FnOnce(usize) -> R) -> R { + self.exports.fetch_add(1, Ordering::Release); + scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); } + // Reading the length also waits for any earlier resize to finish. + f(self.__len__()) + } + + fn buffer_snapshot(buffer: &PyBuffer, vm: &VirtualMachine) -> PyResult> { + let mut bytes = Vec::new(); + bytes + .try_reserve_exact(buffer.desc.len) + .map_err(|_| vm.no_memory_error())?; + buffer.append_to(&mut bytes); + Ok(bytes) + } + + fn search_exported( + &self, + options: ByteInnerFindOptions, + search: impl FnOnce(&[u8], &[u8], core::ops::Range) -> R, + vm: &VirtualMachine, + ) -> PyResult { + self.with_export(|len| { + let (sub, range) = options.parts(len); + let buffer = if sub.check_buffer() { + Some(PyBuffer::from_object(vm, sub, BufferFlags::SIMPLE)?) + } else { + None + }; + let needle = match &buffer { + Some(buffer) => Self::buffer_snapshot(buffer, vm)?, + None => ByteInnerSub::try_from_object(vm, sub.to_owned())?.into_vec(vm)?, + }; + let result = search(&self.borrow_buf(), &needle, range); + drop(buffer); + Ok(result) + }) + } + + fn starts_ends_with( + &self, + options: anystr::StartsEndsWithArgs, + ends: bool, + vm: &VirtualMachine, + ) -> PyResult { + self.with_export(|len| { + let (affix, range) = options.parts(len); + let tailmatch = |item: &PyObject| -> PyResult { + // bytes subclasses use their bytes directly, as in tailmatch. + let buffer = if item.downcast_ref::().is_some() { + None + } else { + Some(PyBuffer::from_object(vm, item, BufferFlags::SIMPLE)?) + }; + let needle = match &buffer { + Some(buffer) => Self::buffer_snapshot(buffer, vm)?, + None => item.downcast_ref::().unwrap().as_bytes().to_vec(), + }; + let result = { + let bytes = self.borrow_buf(); + let slice = match &range { + Some(range) if !range.is_normal() => None, + Some(range) => bytes.get(range.clone()), + None => Some(&*bytes), + }; + slice.is_some_and(|slice| { + if ends { + slice.ends_with(&needle) + } else { + slice.starts_with(&needle) + } + }) + }; + drop(buffer); + Ok(result) + }; + if let Some(tuple) = affix.downcast_ref::() { + for item in tuple { + if tailmatch(item)? { + return Ok(true); + } + } + Ok(false) + } else { + tailmatch(affix).map_err(|error| { + if error.fast_isinstance(vm.ctx.exceptions.type_error) { + vm.new_type_error(format!( + "{} first arg must be bytes or a tuple of bytes, not {}", + if ends { "endswith" } else { "startswith" }, + affix.class().name() + )) + } else { + error + } + }) + } + }) + } + fn repeat(&self, value: isize, vm: &VirtualMachine) -> PyResult { self.inner().mul(value, vm).map(|x| x.into()) } @@ -140,7 +254,7 @@ impl PyByteArray { pub fn _delitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<()> { match SequenceIndex::try_from_borrowed_object::(vm, needle)? { - SequenceIndex::Int(i) => self.try_resizable(vm)?.elements.delitem_by_index(vm, i), + SequenceIndex::Int(i) => self.delitem_by_index(i, vm), SequenceIndex::Slice(slice) => { // TODO: delete 0 elements don't need resizable self.try_resizable(vm)?.elements.delitem_by_slice(vm, slice) @@ -148,6 +262,27 @@ impl PyByteArray { } } + fn delitem_by_index(&self, index: isize, vm: &VirtualMachine) -> PyResult<()> { + // Validate before export rejection when the bytes are available without + // waiting. A codec callback may already hold a read guard with a writer + // queued, so a blocking read here could deadlock before that rejection. + let valid = self + .inner + .try_read() + .map(|inner| inner.elements.wrap_index(index).is_some()); + if valid == Some(false) { + return Err(vm.new_index_error("bytearray index out of range")); + } + let mut inner = self.try_resizable(vm)?; + // Another thread may have resized the bytearray between the locks. + let index = inner + .elements + .wrap_index(index) + .ok_or_else(|| vm.new_index_error("bytearray index out of range"))?; + inner.elements.remove(index); + Ok(()) + } + fn irepeat(zelf: &Py, n: isize, vm: &VirtualMachine) -> PyResult<()> { if n == 1 { return Ok(()); @@ -173,11 +308,6 @@ impl PyByteArray { fn inner(&self) -> PyDetachingRwLockReadGuard<'_, PyBytesInner> { self.inner.read() } - #[inline] - fn inner_mut(&self) -> PyDetachingRwLockWriteGuard<'_, PyBytesInner> { - self.inner.write() - } - fn __len__(&self) -> usize { self.borrow_buf().len() } @@ -496,7 +626,11 @@ impl Py { options: ByteInnerPaddingOptions, vm: &VirtualMachine, ) -> PyResult { - Ok(self.inner().center(options, vm)?.into()) + let (width, fillchar) = options.get_value("center", vm)?; + Ok(self + .inner() + .pad_with(width, fillchar, AnyStr::py_center, vm)? + .into()) } #[pymethod] @@ -505,7 +639,11 @@ impl Py { options: ByteInnerPaddingOptions, vm: &VirtualMachine, ) -> PyResult { - Ok(self.inner().ljust(options, vm)?.into()) + let (width, fillchar) = options.get_value("ljust", vm)?; + Ok(self + .inner() + .pad_with(width, fillchar, AnyStr::py_ljust, vm)? + .into()) } #[pymethod] @@ -514,14 +652,20 @@ impl Py { options: ByteInnerPaddingOptions, vm: &VirtualMachine, ) -> PyResult { - Ok(self.inner().rjust(options, vm)?.into()) + let (width, fillchar) = options.get_value("rjust", vm)?; + Ok(self + .inner() + .pad_with(width, fillchar, AnyStr::py_rjust, vm)? + .into()) } #[pymethod] fn count(&self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult { - self.exports.fetch_add(1, Ordering::Release); - scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); } - self.inner().count(options, vm) + self.search_exported( + options, + |bytes, needle, range| bytes.py_count(needle, range, |h, n| h.find_iter(n).count()), + vm, + ) } #[pymethod] @@ -539,19 +683,7 @@ impl Py { #[pymethod] fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult { - let borrowed = self.borrow_buf(); - let (affix, substr) = - match options.prepare(&*borrowed, borrowed.len(), |s, r| s.get_bytes(r)) { - Some(x) => x, - None => return Ok(false), - }; - substr.py_starts_ends_with( - &affix, - "endswith", - "bytes", - |s, x: PyBytesInner| s.ends_with(x.as_bytes()), - vm, - ) + self.starts_ends_with(options, true, vm) } #[pymethod] @@ -560,19 +692,7 @@ impl Py { options: anystr::StartsEndsWithArgs, vm: &VirtualMachine, ) -> PyResult { - let borrowed = self.borrow_buf(); - let (affix, substr) = - match options.prepare(&*borrowed, borrowed.len(), |s, r| s.get_bytes(r)) { - Some(x) => x, - None => return Ok(false), - }; - substr.py_starts_ends_with( - &affix, - "startswith", - "bytes", - |s, x: PyBytesInner| s.starts_with(x.as_bytes()), - vm, - ) + self.starts_ends_with(options, false, vm) } /// Converting the needle runs Python, which must not resize this @@ -586,9 +706,11 @@ impl Py { where F: Fn(&[u8], &[u8]) -> Option, { - self.exports.fetch_add(1, Ordering::Release); - scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); } - self.inner().find(options, find, vm) + self.search_exported( + options, + |bytes, needle, range| bytes.py_find(needle, range, find), + vm, + ) } #[pymethod] @@ -621,7 +743,28 @@ impl Py { options: ByteInnerTranslateOptions, vm: &VirtualMachine, ) -> PyResult { - Ok(self.inner().translate(options, vm)?.into()) + let mut buffers = ByteArrayBuffers::default(); + if let Some(table) = &options.table { + let buffer = PyBuffer::from_object(vm, table, BufferFlags::SIMPLE)?; + if buffer.desc.len != 256 { + return Err(vm.new_value_error("translation table must be 256 characters long")); + } + buffers.first = Some(buffer); + } + buffers.second = Some(PyBuffer::from_object( + vm, + &options.delete, + BufferFlags::SIMPLE, + )?); + // Both acquisition callbacks finish before any input bytes are read. + let table = match &buffers.first { + Some(buffer) => PyByteArray::buffer_snapshot(buffer, vm)?, + None => (0..=u8::MAX).collect(), + }; + let delete = PyByteArray::buffer_snapshot(buffers.second.as_ref().unwrap(), vm)?; + let result = self.inner().translate_with(&table, &delete).into(); + drop(buffers); + Ok(result) } #[pymethod] @@ -781,9 +924,20 @@ impl Initializer for PyByteArray { type Args = ByteInnerNewOptions; fn init(zelf: &Py, options: Self::Args, vm: &VirtualMachine) -> PyResult<()> { - // First unpack bytearray and *then* get a lock to set it. - let mut inner = options.get_inner(bytearray_from_object, vm)?; - core::mem::swap(&mut *zelf.inner_mut(), &mut inner); + let source_is_missing = matches!(&options.source, OptionalArg::Missing); + // Reject exports and clear existing contents before converting source. + zelf.try_resizable(vm)?.elements = Vec::new(); + // Conversion can call Python, so do not hold the lock across it and + // check again for exports acquired by a callback or another thread. + let (mut inner, zero_count) = + options.get_inner_with_zero_count(bytearray_from_object, vm)?; + // Missing sources and zero integer counts perform no further writes. + // Argument validation above still rejects encoding/errors without a + // source; a zero __index__ callback may have changed or exported self. + if source_is_missing || zero_count { + return Ok(()); + } + core::mem::swap(&mut *zelf.try_resizable(vm)?, &mut inner); Ok(()) } } @@ -795,10 +949,33 @@ impl Comparable for PyByteArray { op: PyComparisonOp, vm: &VirtualMachine, ) -> PyResult { - if let Some(res) = op.identical_optimization(zelf, other) { - return Ok(res.into()); + if !zelf.as_object().check_buffer() || !other.check_buffer() { + return Ok(PyComparisonValue::NotImplemented); } - Ok(zelf.inner().cmp(other, op, vm)) + let first = match PyBuffer::from_object(vm, zelf.as_object(), BufferFlags::SIMPLE) { + Ok(buffer) => buffer, + Err(_) => return Ok(PyComparisonValue::NotImplemented), + }; + let mut buffers = ByteArrayBuffers { + first: Some(first), + second: None, + }; + buffers.second = match PyBuffer::from_object(vm, other, BufferFlags::SIMPLE) { + Ok(buffer) => Some(buffer), + Err(_) => return Ok(PyComparisonValue::NotImplemented), + }; + let left = buffers.first.as_ref().unwrap(); + let right = buffers.second.as_ref().unwrap(); + let result = if left.desc.len != right.desc.len && op.eval_eq(|| false).is_some() { + op.eval_eq(|| false).unwrap() + } else { + let left = Self::buffer_snapshot(left, vm)?; + let right = Self::buffer_snapshot(right, vm)?; + op.eval_eq(|| left == right) + .unwrap_or_else(|| op.eval_ord(left.cmp(&right))) + }; + drop(buffers); + Ok(result.into()) } } @@ -841,9 +1018,11 @@ impl AsBuffer for PyByteArray { } fn as_buffer(zelf: &Py, _vm: &VirtualMachine) -> PyResult { + // Keep the length stable until PyBuffer::new retains the export. + let inner = zelf.inner(); Ok(PyBuffer::new( zelf.to_owned().into(), - BufferDescriptor::simple(zelf.__len__(), false), + BufferDescriptor::simple(inner.elements.len(), false), &BUFFER_METHODS, )) } @@ -855,7 +1034,12 @@ impl BufferResizeGuard for PyByteArray { fn try_resizable_opt(&self) -> Option> { // An export is a borrow someone else still holds, so it is answered // before the lock rather than by waiting on it. - (self.exports.load(Ordering::SeqCst) == 0).then(|| self.inner.write()) + if self.exports.load(Ordering::SeqCst) != 0 { + return None; + } + let inner = self.inner.write(); + // An export may have been acquired while we waited for the lock. + (self.exports.load(Ordering::SeqCst) == 0).then_some(inner) } } @@ -886,10 +1070,43 @@ impl AsSequence for PyByteArray { static AS_SEQUENCE: PySequenceMethods = PySequenceMethods { length: atomic_func!(|seq, _vm| Ok(PyByteArray::sequence_downcast(seq).__len__())), concat: atomic_func!(|seq, other, vm| { - PyByteArray::sequence_downcast(seq) - .inner() - .concat(other, vm) - .map(|x| PyByteArray::from(x).into_pyobject(vm)) + let concat_error = || { + vm.new_type_error(format!( + "can't concat {} to {}", + other.class().name(), + seq.obj.class().name() + )) + }; + // Protect the left operand before the right buffer callback, + // without holding either operand's lock across Python calls. + let left = PyBuffer::from_object(vm, seq.obj, BufferFlags::SIMPLE) + .map_err(|_| concat_error())?; + let right = match PyBuffer::from_object(vm, other, BufferFlags::SIMPLE) { + Ok(buffer) => buffer, + Err(_) => { + let error = concat_error(); + drop(left); + return Err(error); + } + }; + let result = (|| { + let len = left + .desc + .len + .checked_add(right.desc.len) + .ok_or_else(|| vm.no_memory_error())?; + let mut bytes = Vec::new(); + bytes + .try_reserve_exact(len) + .map_err(|_| vm.no_memory_error())?; + left.append_to(&mut bytes); + right.append_to(&mut bytes); + Ok(PyByteArray::from(bytes).into_pyobject(vm)) + })(); + // Buffer release can run Python too; all slice guards are gone. + drop(left); + drop(right); + result }), repeat: atomic_func!(|seq, n, vm| { PyByteArray::sequence_downcast(seq) @@ -907,7 +1124,7 @@ impl AsSequence for PyByteArray { if let Some(value) = value { zelf._setitem_by_index(i, &value, vm) } else { - zelf.borrow_buf_mut().delitem_by_index(vm, i) + zelf.delitem_by_index(i, vm) } }), contains: atomic_func!(|seq, other, vm| { diff --git a/crates/vm/src/builtins/dict.rs b/crates/vm/src/builtins/dict.rs index a98ba4a4356..e59a95d1354 100644 --- a/crates/vm/src/builtins/dict.rs +++ b/crates/vm/src/builtins/dict.rs @@ -7,22 +7,26 @@ use crate::common::lock::LazyLock; use crate::object::{Traverse, TraverseFn}; use crate::stdlib::PyOrderedDictItems; use crate::{ - AsObject, Context, Py, PyExact, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult, + AsObject, Context, Py, PyExact, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, atomic_func, builtins::{PyList, PyTuple, iter::builtins_iter, type_::PyAttributes}, class::{PyClassDef, PyClassImpl}, - common::{ascii, hash::PyHash}, + common::{ + ascii, + hash::{self, PyHash}, + }, dict_inner::{self, DictKey}, function::{ArgIterable, FuncArgs, KwArgs, OptionalArg, PyArithmeticValue, PyComparisonValue}, protocol::{PyIter, PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods}, recursion::ReprGuard, types::{ AsMapping, AsNumber, AsSequence, Callable, Comparable, Constructor, DefaultConstructor, - Initializer, IterNext, Iterable, PyComparisonOp, Representable, SelfIter, + Hashable, Initializer, IterNext, Iterable, PyComparisonOp, Representable, SelfIter, }, vm::VirtualMachine, }; use alloc::fmt; +use core::borrow::Borrow; use core::cell::Cell; use core::ptr::NonNull; use rustpython_common::atomic::{Ordering, PyAtomic, Radium}; @@ -31,15 +35,228 @@ use rustpython_common::wtf8::Wtf8Buf; pub(crate) type DictContentType = dict_inner::Dict; +pub(crate) fn key_hash( + key: &K, + mapping_name: &str, + vm: &VirtualMachine, +) -> PyResult { + match key.key_hash(vm) { + // Exact `TypeError` only: a `__hash__` raising a subclass must + // propagate unchanged, matching CPython's exact-type check. + Err(exc) if exc.class().is(vm.ctx.exceptions.type_error) => { + let name = key.to_pyobject(vm).class().fully_qualified_name(vm)?; + // Match BaseException_str without the diagnostic formatter's + // fallback: an argument's failing __str__/__repr__ must propagate. + let args = exc.args(); + let message = match args.as_slice() { + [] => vm.ctx.empty_str.to_owned(), + [arg] => arg.str(vm)?, + _ => args.as_object().str(vm)?, + }; + Err(vm.new_type_error(format!( + "cannot use '{name}' as a {mapping_name} key ({message})" + ))) + } + result => result, + } +} + #[pyclass(module = false, name = "dict", unhashable = true, traverse = "manual")] #[derive(Default)] // OrderedDict contains eight-byte-aligned fields on 32-bit targets too. #[repr(align(8))] pub struct PyDict { - entries: DictContentType, + pub(crate) entries: DictContentType, } pub type PyDictRef = PyRef; +/// A checked reference to a dict or frozendict, preserving the object's mutability. +/// Its single-pointer layout also supports readonly object members such as __globals__. +#[repr(transparent)] +#[derive(Debug, Clone)] +pub struct PyAnyDictRef(PyObjectRef); + +impl super::descriptor::MemberLayout for PyAnyDictRef { + const KIND: super::descriptor::MemberKind = super::descriptor::MemberKind::Object; +} + +unsafe impl Traverse for PyAnyDictRef { + fn traverse(&self, tracer: &mut TraverseFn<'_>) { + self.0.traverse(tracer); + } +} + +impl Borrow for PyAnyDictRef { + fn borrow(&self) -> &PyObject { + &self.0 + } +} + +impl From for PyAnyDictRef { + fn from(dict: PyDictRef) -> Self { + Self(dict.into()) + } +} +impl From> for PyAnyDictRef { + fn from(dict: PyRef) -> Self { + Self(dict.into()) + } +} +impl From for PyObjectRef { + fn from(dict: PyAnyDictRef) -> Self { + dict.0 + } +} +impl crate::convert::ToPyObject for PyAnyDictRef { + fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef { + self.into() + } +} +impl TryFromObject for PyAnyDictRef { + fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { + Self::from_object(&obj).ok_or_else(|| { + vm.new_type_error(format!( + "expected dict or frozendict, got {}", + obj.class().name() + )) + }) + } +} + +impl PyAnyDictRef { + pub fn from_object(obj: &PyObject) -> Option { + (obj.downcastable::() || obj.downcastable::()) + .then(|| Self(obj.to_owned())) + } + + pub(crate) fn as_mutable(&self) -> Option<&Py> { + self.0.downcast_ref() + } + + pub(crate) fn is_frozen(&self) -> bool { + self.0.downcastable::() + } + + // The raw payload is private: callers outside the VM cannot use it to mutate a frozendict. + pub(crate) fn as_dict(&self) -> &PyDict { + if let Some(dict) = self.as_mutable() { + dict + } else { + &self.0.downcast_ref::().unwrap().dict + } + } + + pub(crate) fn uses_builtin_iter(&self, vm: &VirtualMachine) -> bool { + let base = if self.is_frozen() { + vm.ctx.types.frozendict_type + } else { + vm.ctx.types.dict_type + }; + self.as_object() + .class() + .slots + .iter + .load() + .map(crate::types::fn_addr) + == base.slots.iter.load().map(crate::types::fn_addr) + } + + #[must_use] + pub fn __len__(&self) -> usize { + self.as_dict().__len__() + } + #[must_use] + pub fn is_empty(&self) -> bool { + self.as_dict().is_empty() + } + #[must_use] + pub fn keys_vec(&self) -> Vec { + self.as_dict().keys_vec() + } + #[must_use] + pub fn values_vec(&self) -> Vec { + self.as_dict().values_vec() + } + #[must_use] + pub fn items_vec(&self) -> Vec<(PyObjectRef, PyObjectRef)> { + self.as_dict().items_vec() + } + #[must_use] + pub fn next_entry(&self, position: usize) -> Option<(usize, PyObjectRef, PyObjectRef)> { + self.as_dict().next_entry(position) + } + pub(crate) fn size(&self) -> crate::dict_inner::DictSize { + self.as_dict().size() + } + pub fn inner_getitem_opt( + &self, + key: &K, + vm: &VirtualMachine, + ) -> PyResult> { + if let Some(dict) = self.as_mutable() { + dict.inner_getitem_opt(key, vm) + } else { + self.0 + .downcast_ref::() + .unwrap() + .inner_getitem_opt(key, vm) + } + } + pub fn get_item_opt( + &self, + key: &K, + vm: &VirtualMachine, + ) -> PyResult> { + if let Some(dict) = self.as_mutable() { + return dict.get_item_opt(key, vm); + } + if self.0.class().is(vm.ctx.types.frozendict_type) { + return self.inner_getitem_opt(key, vm); + } + match self.0.get_item(key, vm) { + Ok(value) => Ok(Some(value)), + Err(exc) if exc.fast_isinstance(vm.ctx.exceptions.key_error) => Ok(None), + Err(exc) => Err(exc), + } + } + pub fn get_item(&self, key: &K, vm: &VirtualMachine) -> PyResult { + if let Some(dict) = self.as_mutable() { + dict.get_item(key, vm) + } else { + self.0.get_item(key, vm) + } + } + pub fn contains_key(&self, key: &K, vm: &VirtualMachine) -> bool { + self.inner_getitem_opt(key, vm) + .is_ok_and(|value| value.is_some()) + } + pub fn set_item( + &self, + key: &K, + value: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + match self.as_mutable() { + Some(dict) => dict.set_item(key, value, vm), + None => Err(vm.new_type_error("frozendict object does not support item assignment")), + } + } + pub fn del_item(&self, key: &K, vm: &VirtualMachine) -> PyResult<()> { + match self.as_mutable() { + Some(dict) => dict.del_item(key, vm), + None => Err(vm.new_type_error("frozendict object does not support item deletion")), + } + } +} + +impl<'a> IntoIterator for &'a PyAnyDictRef { + type Item = (PyObjectRef, PyObjectRef); + type IntoIter = DictIter<'a>; + fn into_iter(self) -> Self::IntoIter { + DictIter::new(self.as_dict()) + } +} + // SAFETY: Traverse properly visits all owned PyObjectRefs unsafe impl Traverse for PyDict { fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) { @@ -146,13 +363,11 @@ impl PyDict { override_existing: bool, vm: &VirtualMachine, ) -> PyResult<()> { - let casted: Result, _> = other.downcast_exact(vm); - let other = match casted { - Ok(dict_other) => { - return self.merge_dict(&dict_other, override_existing, vm); - } - Err(other) => other, - }; + if let Some(other_dict) = PyAnyDictRef::from_object(&other) + && other_dict.uses_builtin_iter(vm) + { + return self.merge_dict(other_dict.as_dict(), override_existing, vm); + } let dict = &self.entries; // Use get_attr to properly invoke __getattribute__ for proxy objects let keys_result = other.get_attr(vm.ctx.intern_str("keys"), vm); @@ -280,7 +495,7 @@ impl PyDict { pub(crate) fn merge_dict( &self, - dict_other: &Py, + dict_other: &Self, override_existing: bool, vm: &VirtualMachine, ) -> PyResult<()> { @@ -314,22 +529,7 @@ impl PyDict { /// Hashing up front also means a `__hash__` that fails only intermittently /// is still reported. fn hash_or_unhashable(key: &K, vm: &VirtualMachine) -> PyResult { - match key.key_hash(vm) { - Ok(hash) => Ok(hash), - // Exact `TypeError` only: a `__hash__` raising a *subclass* of - // `TypeError` (e.g. a user `MyTypeError`) must propagate unchanged, - // matching CPython's exact-type check. - Err(cause) if cause.class().is(vm.ctx.exceptions.type_error) => { - let message = cause.as_object().str(vm)?; - let key = key.to_pyobject(vm); - let name = key.class().fully_qualified_name(vm)?; - let err = - vm.new_type_error(format!("cannot use '{name}' as a dict key ({message})")); - err.set_cause(Some(cause)); - Err(err) - } - Err(other) => Err(other), - } + key_hash(key, "dict", vm) } /// Set item variant which can be called with multiple @@ -409,6 +609,8 @@ impl PyDict { ) -> Option> { if let Some(dict) = obj.downcast_ref_if_exact::(vm) { Some(dict.entries.keys_with_hashes()) + } else if let Some(frozen) = obj.downcast_ref_if_exact::(vm) { + Some(frozen.dict.entries.keys_with_hashes()) } else { set::exact_set_keys_with_hashes(obj, vm) } @@ -426,13 +628,48 @@ struct DictGetArgs { #[derive(FromArgs)] struct FromKeysArgs { #[pyarg(positional)] - iterable: ArgIterable, + iterable: PyObjectRef, #[pyarg(positional, optional)] value: Option, } // Python dict methods: impl PyDict { + pub(crate) fn compare_entries( + &self, + other: &Self, + op: PyComparisonOp, + item: bool, + vm: &VirtualMachine, + ) -> PyResult { + if op == PyComparisonOp::Ne { + return Self::compare_entries(self, other, PyComparisonOp::Eq, item, vm) + .map(|x| x.map(|eq| !eq)); + } + if !op.eval_ord(self.__len__().cmp(&other.__len__())) { + return Ok(PyArithmeticValue::Implemented(false)); + } + let mut position = 0; + while let Some((next, k, v1, hash)) = self.entries.next_entry_with_hash(position) { + position = next; + let found = other.entries.get(vm, &*k, hash)?; + match found { + Some(v2) => { + if v1.is(&v2) { + continue; + } + if item && !vm.bool_eq(&v1, &v2)? { + return Ok(PyArithmeticValue::Implemented(false)); + } + } + None => { + return Ok(PyArithmeticValue::Implemented(false)); + } + } + } + Ok(PyArithmeticValue::Implemented(true)) + } + pub fn __len__(&self) -> usize { self.entries.len() } @@ -460,17 +697,6 @@ impl PyDict { self.entries.setdefault(vm, &*key, hash, || default) } - fn __or__(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { - // Accept dict subclasses that inherit this implementation. Subclasses with - // their own reflected slot, such as OrderedDict, take precedence in dispatch. - if let Ok(other) = other.downcast::() { - let self_cp = self.copy(); - self_cp.merge_dict(&other, true, vm)?; - return Ok(self_cp.into_pyobject(vm)); - } - Ok(vm.ctx.not_implemented()) - } - pub(crate) fn __sizeof__(&self) -> usize { core::mem::size_of::() + self.entries.sizeof() } @@ -519,88 +745,20 @@ impl PyDict {} #[pyclass] impl Py { - pub(crate) fn inner_cmp( - &self, - other: &Self, - op: PyComparisonOp, - item: bool, - vm: &VirtualMachine, - ) -> PyResult { - if op == PyComparisonOp::Ne { - return Self::inner_cmp(self, other, PyComparisonOp::Eq, item, vm) - .map(|x| x.map(|eq| !eq)); - } - if !op.eval_ord(self.__len__().cmp(&other.__len__())) { - return Ok(PyArithmeticValue::Implemented(false)); - } - let (superset, subset) = if self.__len__() < other.__len__() { - (other, self) - } else { - (self, other) - }; - for (k, v1) in subset { - match superset.get_item_opt(&*k, vm)? { - Some(v2) => { - if v1.is(&v2) { - continue; - } - if item && !vm.bool_eq(&v1, &v2)? { - return Ok(PyArithmeticValue::Implemented(false)); - } - } - None => { - return Ok(PyArithmeticValue::Implemented(false)); - } - } - } - Ok(PyArithmeticValue::Implemented(true)) - } - #[cfg_attr(feature = "flame-it", flame("PyDictRef"))] #[pymethod(coexist)] fn __getitem__(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult { self.inner_getitem(&*key, vm) } - fn __ror__(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { - let other_dict = other.downcast::(); - if let Ok(other) = other_dict { - let other_cp = other.copy(); - other_cp.merge_dict(self, true, vm)?; - return Ok(other_cp.into_pyobject(vm)); - } - Ok(vm.ctx.not_implemented()) - } - #[pyclassmethod] fn fromkeys(class: PyTypeRef, args: FromKeysArgs, vm: &VirtualMachine) -> PyResult { - let FromKeysArgs { iterable, value } = args; - let value = value.unwrap_or_else(|| vm.ctx.none()); - let d = PyType::call(&class, ().into(), vm)?; - match d.downcast_exact::(vm) { - Ok(pydict) => { - if let Some(keys) = PyDict::fromkeys_known_hashes(iterable.as_object(), vm) { - if class.is(vm.ctx.types.dict_type) { - pydict.entries.reserve_for_empty(keys.len()); - } - for (key, hash) in keys { - pydict.entries.insert(vm, &*key, hash, value.clone())?; - } - } else { - for key in iterable.iter(vm)? { - let key = key?; - pydict.__setitem__(&key, value.clone(), vm)?; - } - } - Ok(pydict.into_pyref().into()) - } - Err(pyobj) => { - for key in iterable.iter(vm)? { - pyobj.set_item(&*key?, value.clone(), vm)?; - } - Ok(pyobj) - } - } + dict_fromkeys( + class, + args.iterable, + args.value.unwrap_or_else(|| vm.ctx.none()), + vm, + ) } #[pymethod] @@ -681,23 +839,23 @@ impl Py { #[pyclass] impl PyRef { #[pymethod] - const fn keys(self) -> PyDictKeys { - PyDictKeys::new(self) + fn keys(self) -> PyDictKeys { + PyDictKeys::new(self.into()) } #[pymethod] - const fn values(self) -> PyDictValues { - PyDictValues::new(self) + fn values(self) -> PyDictValues { + PyDictValues::new(self.into()) } #[pymethod] - const fn items(self) -> PyDictItems { - PyDictItems::new(self) + fn items(self) -> PyDictItems { + PyDictItems::new(self.into()) } #[pymethod] fn __reversed__(self) -> PyDictReverseKeyIterator { - PyDictReverseKeyIterator::new(self) + PyDictReverseKeyIterator::new(self.into()) } fn __ior__(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { @@ -758,13 +916,7 @@ impl AsSequence for PyDict { impl AsNumber for PyDict { fn as_number() -> &'static PyNumberMethods { static AS_NUMBER: PyNumberMethods = PyNumberMethods { - or: Some(|a, b, vm| { - if let Some(a) = a.downcast_ref::() { - PyDict::__or__(a, b.to_pyobject(vm), vm) - } else { - Ok(vm.ctx.not_implemented()) - } - }), + or: Some(PyFrozenDict::union), inplace_or: Some(|a, b, vm| { if let Some(a) = a.downcast_ref::() { a.to_owned() @@ -787,16 +939,19 @@ impl Comparable for PyDict { op: PyComparisonOp, vm: &VirtualMachine, ) -> PyResult { - op.eq_only(|| { - let other = class_or_notimplemented!(Self, other); - zelf.inner_cmp(other, PyComparisonOp::Eq, true, vm) + op.eq_only(|| match PyAnyDictRef::from_object(other) { + Some(other) => { + zelf.payload + .compare_entries(other.as_dict(), PyComparisonOp::Eq, true, vm) + } + None => Ok(PyComparisonValue::NotImplemented), }) } } impl Iterable for PyDict { fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { - Ok(PyDictKeyIterator::new(zelf).into_pyobject(vm)) + Ok(PyDictKeyIterator::new(zelf.into()).into_pyobject(vm)) } } @@ -1165,6 +1320,387 @@ impl PyExact { } } +#[doc = r"frozendict() -> new empty immutable dictionary +frozendict(mapping) -> new immutable dictionary initialized from a mapping + object's (key, value) pairs +frozendict(iterable) -> new immutable dictionary initialized as if via: + d = {} + for k, v in iterable: + d[k] = v + d = frozendict(d) +frozendict(**kwargs) -> new immutable dictionary initialized with the name=value + pairs in the keyword argument list. For example: frozendict(one=1, two=2)"] +#[pyclass(module = false, name = "frozendict", traverse = "manual")] +#[derive(Debug)] +pub struct PyFrozenDict { + // Own the table directly, never a mutable Python dict visible through gc or a view. + pub(crate) dict: PyDict, + hash: PyAtomic, +} + +unsafe impl Traverse for PyFrozenDict { + fn traverse(&self, tracer: &mut TraverseFn<'_>) { + self.dict.traverse(tracer); + } + fn clear(&mut self, out: &mut Vec) { + Traverse::clear(&mut self.dict, out); + } +} + +impl PyPayload for PyFrozenDict { + fn class(ctx: &Context) -> &'static Py { + ctx.types.frozendict_type + } +} + +impl PyFrozenDict { + pub(crate) fn from_dict(dict: PyDict) -> Self { + Self { + dict, + hash: Radium::new(hash::SENTINEL), + } + } + + pub fn from_object(input: Option, vm: &VirtualMachine) -> PyResult> { + if let Some(input) = &input + && let Some(frozen) = input.downcast_ref_if_exact::(vm) + { + return Ok(frozen.to_owned()); + } + let dict = PyDict::default(); + if let Some(input) = input { + Self::merge_into(&dict, input, vm)?; + } + Ok(Self::from_dict(dict).into_ref(&vm.ctx)) + } + + /// Hash a dictionary key, adding the frozendict context to an exact TypeError. + pub fn key_hash(key: &K, vm: &VirtualMachine) -> PyResult { + key_hash(key, "frozendict", vm) + } + + pub fn inner_getitem_opt( + &self, + key: &K, + vm: &VirtualMachine, + ) -> PyResult> { + let hash = Self::key_hash(key, vm)?; + self.dict.entries.get(vm, key, hash) + } + + fn insert( + dict: &PyDict, + key: &PyObject, + value: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + let hash = Self::key_hash(key, vm)?; + dict.entries.insert(vm, key, hash, value) + } + + // This table is unpublished until construction/union/fromkeys finishes. Python callbacks + // cannot find and mutate a partially initialized frozendict through the GC. + pub(crate) fn merge_into( + dict: &PyDict, + other: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + if let Some(other_dict) = PyAnyDictRef::from_object(&other) + && other_dict.uses_builtin_iter(vm) + { + return dict.merge_dict(other_dict.as_dict(), true, vm); + } + match other.get_attr(vm.ctx.intern_str("keys"), vm) { + Ok(keys) => { + let keys = PyIter::try_from_object(vm, keys.call((), vm)?)?; + while let PyIterReturn::Return(key) = keys.next(vm)? { + let value = other.get_item(&*key, vm)?; + Self::insert(dict, &key, value, vm)?; + } + } + Err(exc) if exc.fast_isinstance(vm.ctx.exceptions.attribute_error) => { + for (index, pair) in other.get_iter(vm)?.iter::(vm)?.enumerate() { + let (key, value) = PyDict::update_sequence_pair(pair?, index, vm)?; + Self::insert(dict, &key, value, vm)?; + } + } + Err(exc) => return Err(exc), + } + Ok(()) + } + + pub(crate) fn union(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult { + let (Some(left), Some(right)) = + (PyAnyDictRef::from_object(a), PyAnyDictRef::from_object(b)) + else { + return Ok(vm.ctx.not_implemented()); + }; + if a.class().is(vm.ctx.types.frozendict_type) { + if left.is_empty() && b.class().is(vm.ctx.types.frozendict_type) { + return Ok(b.to_owned()); + } + if right.is_empty() + && (b.class().is(vm.ctx.types.dict_type) + || b.class().is(vm.ctx.types.frozendict_type)) + { + return Ok(a.to_owned()); + } + } + let dict = PyDict::default(); + if left.is_frozen() { + if !left.is_empty() { + Self::merge_into(&dict, a.to_owned(), vm)?; + } + Self::merge_into(&dict, b.to_owned(), vm)?; + Ok(Self::from_dict(dict).into_pyobject(vm)) + } else { + if !left.is_empty() { + dict.merge_object(a.to_owned(), vm)?; + } + dict.merge_object(b.to_owned(), vm)?; + Ok(dict.into_pyobject(vm)) + } + } +} + +impl Constructor for PyFrozenDict { + type Args = (OptionalArg, KwArgs); + fn py_new( + _cls: &Py, + (input, kwargs): Self::Args, + vm: &VirtualMachine, + ) -> PyResult { + let dict = PyDict::default(); + if let OptionalArg::Present(input) = input { + Self::merge_into(&dict, input, vm)?; + } + for (key, value) in kwargs { + dict.inner_setitem(&key, value, vm)?; + } + Ok(Self::from_dict(dict)) + } +} + +#[pyclass( + with( + Py, + PyRef, + Constructor, + Hashable, + Comparable, + Iterable, + AsMapping, + AsSequence, + AsNumber, + Representable + ), + flags(BASETYPE, MAPPING, _MATCH_SELF) +)] +impl PyFrozenDict { + fn __len__(&self) -> usize { + self.dict.__len__() + } + + #[pymethod] + fn get( + zelf: &Py, + key: PyObjectRef, + default: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + Ok(zelf + .inner_getitem_opt(&*key, vm)? + .unwrap_or_else(|| default.unwrap_or_else(|| vm.ctx.none()))) + } + + #[pymethod] + fn __sizeof__(zelf: &Py) -> usize { + zelf.dict.__sizeof__() + core::mem::size_of::>() + } + + #[pyclassmethod] + fn fromkeys( + cls: PyTypeRef, + iterable: PyObjectRef, + value: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + dict_fromkeys(cls, iterable, value.unwrap_or_else(|| vm.ctx.none()), vm) + } + + #[pyclassmethod] + fn __class_getitem__( + cls: PyTypeRef, + object: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + PyGenericAlias::from_args(cls, object, vm) + } +} + +#[pyclass] +impl Py { + #[pymethod(coexist)] + fn __getitem__(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult { + if let Some(value) = self.inner_getitem_opt(&*key, vm)? { + return Ok(value); + } + if !self.class().is(vm.ctx.types.frozendict_type) + && let Some(method) = + vm.get_method(self.to_owned().into(), identifier!(vm, __missing__)) + { + return method?.call((key,), vm); + } + Err(vm.new_key_error(key)) + } + + #[pymethod] + fn copy(&self, vm: &VirtualMachine) -> PyResult> { + if self.class().is(vm.ctx.types.frozendict_type) { + Ok(self.to_owned()) + } else { + let dict = PyDict::default(); + if !self.dict.is_empty() { + PyFrozenDict::merge_into(&dict, self.to_owned().into(), vm)?; + } + Ok(PyFrozenDict::from_dict(dict).into_ref(&vm.ctx)) + } + } + + #[pymethod] + fn __getnewargs__(&self, vm: &VirtualMachine) -> PyResult { + let dict = vm + .ctx + .types + .dict_type + .as_object() + .call((self.to_owned(),), vm)?; + Ok(vm.new_tuple((dict,))) + } +} + +#[pyclass] +impl PyRef { + #[pymethod] + fn keys(self) -> PyDictKeys { + PyDictKeys::new(self.into()) + } + #[pymethod] + fn values(self) -> PyDictValues { + PyDictValues::new(self.into()) + } + #[pymethod] + fn items(self) -> PyDictItems { + PyDictItems::new(self.into()) + } + #[pymethod] + fn __reversed__(self) -> PyDictReverseKeyIterator { + PyDictReverseKeyIterator::new(self.into()) + } +} + +impl AsMapping for PyFrozenDict { + fn as_mapping() -> &'static PyMappingMethods { + static METHODS: PyMappingMethods = PyMappingMethods { + length: atomic_func!(|obj, _vm| Ok(PyFrozenDict::mapping_downcast(obj).__len__())), + subscript: atomic_func!( + |obj, key, vm| PyFrozenDict::mapping_downcast(obj).__getitem__(key.to_owned(), vm) + ), + ass_subscript: None, + }; + &METHODS + } +} +impl AsSequence for PyFrozenDict { + fn as_sequence() -> &'static PySequenceMethods { + static METHODS: LazyLock = LazyLock::new(|| PySequenceMethods { + contains: atomic_func!(|obj, key, vm| Ok(PyFrozenDict::sequence_downcast(obj) + .inner_getitem_opt(key, vm)? + .is_some())), + ..PySequenceMethods::NOT_IMPLEMENTED + }); + &METHODS + } +} +impl AsNumber for PyFrozenDict { + fn as_number() -> &'static PyNumberMethods { + static METHODS: PyNumberMethods = PyNumberMethods { + or: Some(PyFrozenDict::union), + ..PyNumberMethods::NOT_IMPLEMENTED + }; + &METHODS + } +} +impl Iterable for PyFrozenDict { + fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + Ok(PyDictKeyIterator::new(zelf.into()).into_pyobject(vm)) + } +} +impl Comparable for PyFrozenDict { + fn cmp( + zelf: &Py, + other: &PyObject, + op: PyComparisonOp, + vm: &VirtualMachine, + ) -> PyResult { + op.eq_only(|| match PyAnyDictRef::from_object(other) { + Some(other) => zelf + .dict + .compare_entries(other.as_dict(), PyComparisonOp::Eq, true, vm), + None => Ok(PyComparisonValue::NotImplemented), + }) + } +} +impl Hashable for PyFrozenDict { + fn hash(zelf: &Py, vm: &VirtualMachine) -> PyResult { + let cached = zelf.hash.load(Ordering::Relaxed); + if cached != hash::SENTINEL { + return Ok(cached); + } + let mut hasher = hash::FrozenSetHash::new(zelf.__len__()); + let mut position = 0; + while let Some((next, _key, value, key_hash)) = + zelf.dict.entries.next_entry_with_hash(position) + { + position = next; + // next_entry releases the table lock before invoking arbitrary Python __hash__. + let pair_hash = hash::hash_tuple([Ok(key_hash), value.hash(vm)])?; + hasher.add(pair_hash); + } + let result = hasher.finish(); + zelf.hash.store(result, Ordering::Relaxed); + Ok(result) + } +} +impl Representable for PyFrozenDict { + fn repr(zelf: &Py, vm: &VirtualMachine) -> PyResult { + let name = zelf.class().name(); + if zelf.dict.is_empty() { + return Ok(vm.ctx.new_str(format!("{name}()"))); + } + let mut result = Wtf8Buf::from(format!("{name}(")); + if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) { + result.push_str("{"); + for (index, (key, value)) in DictIter::new(&zelf.dict).enumerate() { + if index != 0 { + result.push_str(", "); + } + result.push_wtf8(key.repr(vm)?.as_wtf8()); + result.push_str(": "); + result.push_wtf8(value.repr(vm)?.as_wtf8()); + } + result.push_str("}"); + } else { + result.push_str("{...}"); + } + result.push_str(")"); + Ok(vm.ctx.new_str(result)) + } + fn repr_str(_zelf: &Py, _vm: &VirtualMachine) -> PyResult { + unreachable!("use repr instead") + } +} + // Implement IntoIterator so that we can easily iterate dictionaries from rust code. impl IntoIterator for PyDictRef { type Item = (PyObjectRef, PyObjectRef); @@ -1226,12 +1762,12 @@ impl ExactSizeIterator for DictIntoIter { } pub struct DictIter<'a> { - dict: &'a Py, + dict: &'a PyDict, position: usize, } impl<'a> DictIter<'a> { - pub const fn new(dict: &'a Py) -> Self { + pub const fn new(dict: &'a PyDict) -> Self { DictIter { dict, position: 0 } } } @@ -1261,7 +1797,7 @@ impl ExactSizeIterator for DictIter<'_> { trait DictView: PyPayload + PyClassDef + Iterable + Representable { type ReverseIter: PyPayload + core::fmt::Debug; - fn dict(&self) -> &Py; + fn dict(&self) -> &PyAnyDictRef; fn item(vm: &VirtualMachine, key: PyObjectRef, value: PyObjectRef) -> PyObjectRef; fn __len__(&self) -> usize { @@ -1287,14 +1823,14 @@ macro_rules! dict_view { $project_fn: expr, $result_fn: expr ) => { - #[pyclass(module = false, name = $class_name, unhashable = $unhashable)] + #[pyclass(module = false, name = $class_name, unhashable = $unhashable, traverse)] #[derive(Debug)] pub(crate) struct $name { - pub(crate) dict: PyDictRef, + pub(crate) dict: PyAnyDictRef, } impl $name { - pub(crate) const fn new(dict: PyDictRef) -> Self { + pub(crate) fn new(dict: PyAnyDictRef) -> Self { $name { dict } } } @@ -1302,7 +1838,7 @@ macro_rules! dict_view { impl DictView for $name { type ReverseIter = $reverse_iter_name; - fn dict(&self) -> &Py { + fn dict(&self) -> &PyAnyDictRef { &self.dict } @@ -1333,7 +1869,7 @@ macro_rules! dict_view { let s = if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) { let mut result = Wtf8Buf::from(format!("{}([", Self::NAME)); let mut first = true; - for (key, value) in zelf.dict().clone() { + for (key, value) in zelf.dict() { if !first { result.push_str(", "); } @@ -1354,16 +1890,18 @@ macro_rules! dict_view { } } - #[pyclass(module = false, name = $iter_class_name)] + #[pyclass(module = false, name = $iter_class_name, traverse)] #[derive(Debug)] pub(crate) struct $iter_name { + #[pytraverse(skip)] pub(crate) size: dict_inner::DictSize, /// Whether the dict was found to have changed, which /// `dictiter_iternextkey()` records by writing a size no dict can /// have. Sticky: what it makes the iterator answer, it answers /// from then on. + #[pytraverse(skip)] changed: PyAtomic, - pub(crate) internal: PyMutex>, + pub(crate) internal: PyMutex>, } impl PyPayload for $iter_name { @@ -1374,7 +1912,7 @@ macro_rules! dict_view { } impl $iter_name { - fn new(dict: PyDictRef) -> Self { + pub(crate) fn new(dict: PyAnyDictRef) -> Self { $iter_name { size: dict.size(), changed: Radium::new(false), @@ -1411,7 +1949,7 @@ macro_rules! dict_view { let mut position = internal.position; let mut entries = Vec::new(); while let Some((next_position, key, value)) = - dict.entries.next_entry(position) + dict.as_dict().entries.next_entry(position) { entries.push(($result_fn)(vm, ($project_fn)(&key, &value))); position = next_position; @@ -1440,9 +1978,11 @@ macro_rules! dict_view { // raising. return (Err(mutated()), None); } - let entry = - dict.entries - .next_entry_checked(internal.position, &zelf.size, $project_fn); + let entry = dict.as_dict().entries.next_entry_checked( + internal.position, + &zelf.size, + $project_fn, + ); match entry { Err(dict_inner::DictChanged) => { zelf.changed.store(true, Ordering::Relaxed); @@ -1458,14 +1998,17 @@ macro_rules! dict_view { } } - #[pyclass(module = false, name = $reverse_iter_class_name)] + #[pyclass(module = false, name = $reverse_iter_class_name, traverse)] #[derive(Debug)] pub(crate) struct $reverse_iter_name { + #[pytraverse(skip)] pub(crate) size: dict_inner::DictSize, /// As in `$iter_name`. + #[pytraverse(skip)] changed: PyAtomic, + #[pytraverse(skip)] remaining: PyAtomic, - internal: PyMutex>, + internal: PyMutex>, } impl PyPayload for $reverse_iter_name { @@ -1476,7 +2019,7 @@ macro_rules! dict_view { } impl $reverse_iter_name { - fn new(dict: PyDictRef) -> Self { + pub(crate) fn new(dict: PyAnyDictRef) -> Self { let size = dict.size(); // usize::MAX represents the position before the first entry. let position = size.entries_size.wrapping_sub(1); @@ -1500,7 +2043,7 @@ macro_rules! dict_view { let mut position = internal.position; let mut entries = Vec::new(); while let Some((found_index, key, value)) = - dict.entries.prev_entry(position) + dict.as_dict().entries.prev_entry(position) { entries.push(($result_fn)(vm, ($project_fn)(&key, &value))); if found_index == 0 { @@ -1549,9 +2092,11 @@ macro_rules! dict_view { // raising. return (Err(mutated()), None); } - let entry = - dict.entries - .prev_entry_checked(internal.position, &zelf.size, $project_fn); + let entry = dict.as_dict().entries.prev_entry_checked( + internal.position, + &zelf.size, + $project_fn, + ); match entry { Err(dict_inner::DictChanged) => { zelf.changed.store(true, Ordering::Relaxed); @@ -1637,14 +2182,6 @@ trait ViewSetOps: DictView { op: PyComparisonOp, vm: &VirtualMachine, ) -> PyResult { - if let Some(dictview) = other.downcast_ref::() { - return zelf.dict().inner_cmp( - dictview.dict(), - op, - !zelf.class().is(vm.ctx.types.dict_keys_type), - vm, - ); - } if !is_set_or_dict_view(other) { return Ok(PyComparisonValue::NotImplemented); } @@ -1718,7 +2255,7 @@ impl PyDictKeys { #[pygetset] fn mapping(zelf: PyRef) -> PyMappingProxy { - PyMappingProxy::from(zelf.dict().to_owned()) + PyMappingProxy::from_any_dict(zelf.dict().clone()) } } @@ -1738,11 +2275,10 @@ impl AsSequence for PyDictKeys { static AS_SEQUENCE: LazyLock = LazyLock::new(|| PySequenceMethods { length: atomic_func!(|seq, _vm| Ok(PyDictKeys::sequence_downcast(seq).__len__())), contains: atomic_func!(|seq, target, vm| { - let hash = PyDict::hash_or_unhashable(target, vm)?; - PyDictKeys::sequence_downcast(seq) + Ok(PyDictKeys::sequence_downcast(seq) .dict - .entries - .contains(vm, target, hash) + .inner_getitem_opt(target, vm)? + .is_some()) }), ..PySequenceMethods::NOT_IMPLEMENTED }); @@ -1782,7 +2318,7 @@ impl PyDictItems { } #[pygetset] fn mapping(zelf: PyRef) -> PyMappingProxy { - PyMappingProxy::from(zelf.dict().to_owned()) + PyMappingProxy::from_any_dict(zelf.dict().clone()) } } @@ -1844,7 +2380,7 @@ impl AsNumber for PyDictItems { impl PyDictValues { #[pygetset] fn mapping(zelf: PyRef) -> PyMappingProxy { - PyMappingProxy::from(zelf.dict().to_owned()) + PyMappingProxy::from_any_dict(zelf.dict().clone()) } } @@ -1944,6 +2480,68 @@ pub(crate) fn set_inner_number_or(a: &PyObject, b: &PyObject, vm: &VirtualMachin set_inner_number_op(a, b, |a, b| a.union(b, vm), vm) } +pub(crate) fn dict_fromkeys( + class: PyTypeRef, + iterable: PyObjectRef, + value: PyObjectRef, + vm: &VirtualMachine, +) -> PyResult { + let exact_frozen = class.is(vm.ctx.types.frozendict_type); + let created = if exact_frozen { + None + } else { + Some(PyType::call(&class, ().into(), vm)?) + }; + let frozen_source = created + .as_ref() + .and_then(|obj| obj.downcast_ref::()); + if exact_frozen || frozen_source.is_some() { + let dict = PyDict::default(); + if let Some(source) = frozen_source { + PyFrozenDict::merge_into(&dict, source.to_owned().into(), vm)?; + } + let iterable = ArgIterable::::try_from_object(vm, iterable)?; + if let Some(keys) = PyDict::fromkeys_known_hashes(iterable.as_object(), vm) { + for (key, hash) in keys { + dict.entries.insert(vm, &*key, hash, value.clone())?; + } + } else { + for key in iterable.iter(vm)? { + let key = key?; + let hash = PyFrozenDict::key_hash(&*key, vm)?; + dict.entries.insert(vm, &*key, hash, value.clone())?; + } + } + let frozen = PyFrozenDict::from_dict(dict).into_ref(&vm.ctx); + return if exact_frozen { + Ok(frozen.into()) + } else { + class.as_object().call((frozen,), vm) + }; + } + let iterable = ArgIterable::::try_from_object(vm, iterable)?; + let created = created.unwrap(); + if let Some(dict) = created.downcast_ref_if_exact::(vm) { + if let Some(keys) = PyDict::fromkeys_known_hashes(iterable.as_object(), vm) { + if class.is(vm.ctx.types.dict_type) { + dict.entries.reserve_for_empty(keys.len()); + } + for (key, hash) in keys { + dict.entries.insert(vm, &*key, hash, value.clone())?; + } + } else { + for key in iterable.iter(vm)? { + dict.inner_setitem(&*key?, value.clone(), vm)?; + } + } + } else { + for key in iterable.iter(vm)? { + created.set_item(&*key?, value.clone(), vm)?; + } + } + Ok(created) +} + fn vectorcall_dict( zelf_obj: &PyObject, args: Vec, @@ -1958,6 +2556,26 @@ fn vectorcall_dict( Ok(obj.into()) } +fn vectorcall_frozendict( + zelf: &PyObject, + args: Vec, + nargs: usize, + kwnames: Option<&[PyObjectRef]>, + vm: &VirtualMachine, +) -> PyResult { + if nargs == 1 + && kwnames.is_none_or(<[_]>::is_empty) + && args[0].class().is(vm.ctx.types.frozendict_type) + { + return Ok(args[0].clone()); + } + PyFrozenDict::slot_new( + zelf.downcast_ref::().unwrap().to_owned(), + FuncArgs::from_vectorcall_owned(args, nargs, kwnames), + vm, + ) +} + pub(crate) fn init(context: &'static Context) { PyDict::extend_class(context, context.types.dict_type); context @@ -1975,4 +2593,11 @@ pub(crate) fn init(context: &'static Context) { PyDictItems::extend_class(context, context.types.dict_items_type); PyDictItemIterator::extend_class(context, context.types.dict_itemiterator_type); PyDictReverseItemIterator::extend_class(context, context.types.dict_reverseitemiterator_type); + PyFrozenDict::extend_class(context, context.types.frozendict_type); + context + .types + .frozendict_type + .slots + .vectorcall + .store(Some(vectorcall_frozendict)); } diff --git a/crates/vm/src/builtins/frame.rs b/crates/vm/src/builtins/frame.rs index 905cbb582a4..49e1d4ad5c0 100644 --- a/crates/vm/src/builtins/frame.rs +++ b/crates/vm/src/builtins/frame.rs @@ -2,7 +2,7 @@ */ -use super::{PyAsyncGen, PyCode, PyCoroutine, PyDictRef, PyGenerator, PyIntRef}; +use super::{PyAnyDictRef, PyAsyncGen, PyCode, PyCoroutine, PyGenerator, PyIntRef}; use crate::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, class::PyClassImpl, @@ -668,7 +668,7 @@ impl Py { } #[pygetset] - pub fn f_globals(&self) -> PyDictRef { + pub fn f_globals(&self) -> PyAnyDictRef { self.iframe().globals().to_owned() } diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index b0bd5a010b5..ff14ef78395 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -2,8 +2,8 @@ mod jit; use super::{ - PyAsyncGen, PyCode, PyCoroutine, PyDictRef, PyGenerator, PyList, PyModule, PyStr, PyStrRef, - PyTuple, PyTupleRef, PyType, object, + PyAnyDictRef, PyAsyncGen, PyCode, PyCoroutine, PyDictRef, PyGenerator, PyList, PyModule, PyStr, + PyStrRef, PyTuple, PyTupleRef, PyType, object, }; use crate::common::hash::PyHash; use crate::common::lock::PyMutex; @@ -66,7 +66,7 @@ fn format_missing_args( pub struct PyFunction { pub(crate) code: PyAtomicRef, #[pymember(name = "__globals__")] - pub(crate) globals: PyDictRef, + pub(crate) globals: PyAnyDictRef, #[pymember(name = "__builtins__")] pub(crate) builtins: PyObjectRef, #[pymember(name = "__closure__")] @@ -179,15 +179,17 @@ impl PyFunction { #[inline] pub(crate) fn new( code: PyRef, - globals: PyDictRef, + globals: PyAnyDictRef, vm: &VirtualMachine, ) -> PyResult { let name = PyMutex::new(code.obj_name.to_owned()); - let module = vm.unwrap_or_none(globals.get_item_opt(identifier!(vm, __name__), vm)?); - let builtins = globals.get_item("__builtins__", vm).unwrap_or_else(|_| { - // If not in globals, inherit from current execution context - crate::frame::current_builtins().unwrap_or_else(|| vm.builtins.dict().into()) - }); + let module = vm.unwrap_or_none(globals.inner_getitem_opt(identifier!(vm, __name__), vm)?); + let builtins = globals + .inner_getitem_opt(identifier!(vm, __builtins__), vm)? + .unwrap_or_else(|| { + // If not in globals, inherit from current execution context. + crate::frame::current_builtins().unwrap_or_else(|| vm.builtins.dict().into()) + }); // If builtins is a module, use its __dict__ instead let builtins = if let Some(module) = builtins.downcast_ref::() { module.dict().into() @@ -621,7 +623,7 @@ impl Py { } else if let Some(locals) = locals { Some(locals) } else { - Some(ArgMapping::from_dict_exact(self.globals.clone())) + Some(ArgMapping::from_anydict_exact(self.globals.clone())) }; let use_datastack = !is_gen && !is_coro && !is_async_gen; let frame = FrameObject::new_ref( @@ -654,7 +656,7 @@ impl Py { } else if let Some(locals) = locals { crate::frame::FrameLocals::with_locals(locals) } else { - crate::frame::FrameLocals::with_locals(crate::function::ArgMapping::from_dict_exact( + crate::frame::FrameLocals::with_locals(crate::function::ArgMapping::from_anydict_exact( self.globals.clone(), )) }; @@ -805,7 +807,7 @@ impl Py { let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) { None } else { - Some(ArgMapping::from_dict_exact(self.globals.clone())) + Some(ArgMapping::from_anydict_exact(self.globals.clone())) }; let frame = FrameObject::new_ref( @@ -862,7 +864,7 @@ impl Py { let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) { None } else { - Some(ArgMapping::from_dict_exact(self.globals.clone())) + Some(ArgMapping::from_anydict_exact(self.globals.clone())) }; // Heap-backed: the frame outlives the call that made it. @@ -897,7 +899,7 @@ impl Py { let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) { crate::frame::FrameLocals::lazy() } else { - crate::frame::FrameLocals::with_locals(ArgMapping::from_dict_exact( + crate::frame::FrameLocals::with_locals(ArgMapping::from_anydict_exact( self.globals.clone(), )) }; @@ -1318,7 +1320,7 @@ impl Constructor for PyFunction { None }; - let mut func = Self::new(args.code.clone(), args.globals.clone(), vm)?; + let mut func = Self::new(args.code.clone(), args.globals.clone().into(), vm)?; // Set function name if provided if let Some(name) = args.name.into_option() { *func.name.lock() = name; diff --git a/crates/vm/src/builtins/mappingproxy.rs b/crates/vm/src/builtins/mappingproxy.rs index 3e53f1729a5..0c302de8229 100644 --- a/crates/vm/src/builtins/mappingproxy.rs +++ b/crates/vm/src/builtins/mappingproxy.rs @@ -68,6 +68,12 @@ impl Constructor for PyMappingProxy { } impl PyMappingProxy { + pub(crate) fn from_any_dict(dict: super::PyAnyDictRef) -> Self { + Self { + mapping: MappingProxyInner::Mapping(ArgMapping::new(dict.into())), + } + } + pub fn from_object(mapping: PyObjectRef, vm: &VirtualMachine) -> PyResult { if mapping.mapping_unchecked().check() && !mapping.downcastable::() diff --git a/crates/vm/src/builtins/mod.rs b/crates/vm/src/builtins/mod.rs index a938b9b3dd1..a31ba4ffa31 100644 --- a/crates/vm/src/builtins/mod.rs +++ b/crates/vm/src/builtins/mod.rs @@ -20,7 +20,7 @@ pub use complex::PyComplex; pub(crate) mod coroutine; pub use coroutine::PyCoroutine; pub(crate) mod dict; -pub use dict::{PyDict, PyDictRef}; +pub use dict::{PyAnyDictRef, PyDict, PyDictRef, PyFrozenDict}; pub(crate) mod enumerate; pub use enumerate::PyEnumerate; pub(crate) mod filter; diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs index a2449f136ef..be75baf5593 100644 --- a/crates/vm/src/builtins/type.rs +++ b/crates/vm/src/builtins/type.rs @@ -1,6 +1,6 @@ use super::{ - PyClassMethod, PyDict, PyDictRef, PyList, PyStaticMethod, PyStr, PyStrInterned, PyStrRef, - PyTupleRef, PyUtf8StrRef, PyWeak, mappingproxy::PyMappingProxy, object, union_, + PyAnyDictRef, PyClassMethod, PyDict, PyDictRef, PyList, PyStaticMethod, PyStr, PyStrInterned, + PyStrRef, PyTupleRef, PyUtf8StrRef, PyWeak, mappingproxy::PyMappingProxy, object, union_, }; use crate::{ AsObject, Context, Py, PyAtomicRef, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, @@ -1798,25 +1798,9 @@ impl Constructor for PyType { })); } - let (name, bases, dict, kwargs): (PyStrRef, PyTupleRef, PyDictRef, KwArgs) = + let (name, bases, dict, kwargs): (PyStrRef, PyTupleRef, PyAnyDictRef, KwArgs) = args.clone().bind_for(vm, Self::NAME)?; - // A mapping that is not an exact dict (e.g. OrderedDict) keeps its - // own iteration order; copy via the mapping protocol so that order - // lands in the type dict. - let dict = if args.args[2].class().is(vm.ctx.types.dict_type) { - dict - } else { - let copied = vm.ctx.new_dict(); - copied.merge_object(args.args[2].clone(), vm)?; - copied - }; - - if name.as_bytes().contains(&0) { - return Err(vm.new_value_error("type name must not contain null characters")); - } - let name = name.try_into_utf8(vm)?; - let (metatype, base, bases, base_is_type) = if bases.as_slice().is_empty() { let base = vm.ctx.types.object_type.to_owned(); let bases = PyTuple::new_ref_typed(vec![base.clone()], &vm.ctx); @@ -1856,6 +1840,25 @@ impl Constructor for PyType { (metatype, base.to_owned(), bases, base_is_type) }; + // Copy only after resolving the metaclass: its __new__ receives the + // original namespace. A subclass with a custom iterator (e.g. + // OrderedDict) keeps its own mapping order in the mutable type dict. + let dict = if let Some(dict) = dict.as_object().downcast_ref_if_exact::(vm) { + dict.to_owned() + } else { + let copied = vm.ctx.new_dict(); + // Like dict.copy(), an empty table bypasses subclass hooks. + if !dict.is_empty() { + copied.merge_object(dict.into(), vm)?; + } + copied + }; + + if name.as_bytes().contains(&0) { + return Err(vm.new_value_error("type name must not contain null characters")); + } + let name = name.try_into_utf8(vm)?; + let qualname = dict .get_item_opt(identifier!(vm, __qualname__), vm)? .map(|obj| downcast_qualname(obj, vm)) diff --git a/crates/vm/src/bytes_inner.rs b/crates/vm/src/bytes_inner.rs index f577e302085..8c52948c831 100644 --- a/crates/vm/src/bytes_inner.rs +++ b/crates/vm/src/bytes_inner.rs @@ -225,12 +225,24 @@ impl ByteInnerNewOptions { /// read: `bytes()` and `bytearray()` differ in whether they ask it how long /// it is. pub fn get_inner(self, from_object: FromObject, vm: &VirtualMachine) -> PyResult { - match (self.source, self.encoding, self.errors) { + self.get_inner_with_zero_count(from_object, vm) + .map(|(inner, _)| inner) + } + + pub(crate) fn get_inner_with_zero_count( + self, + from_object: FromObject, + vm: &VirtualMachine, + ) -> PyResult<(PyBytesInner, bool)> { + let mut zero_count = false; + let inner = match (self.source, self.encoding, self.errors) { (OptionalArg::Present(obj), OptionalArg::Missing, OptionalArg::Missing) => { // Try __index__ first to handle int-like objects that might raise custom exceptions if let Some(index_result) = obj.try_index_opt(vm) { match index_result { - Ok(index) => Self::get_value_from_size(&index, vm), + Ok(index) => Self::get_value_from_size(&index, vm).inspect(|inner| { + zero_count = inner.elements.is_empty(); + }), Err(e) => { // Only propagate non-TypeError exceptions // TypeError means the object doesn't support __index__, so fall back @@ -266,7 +278,8 @@ impl ByteInnerNewOptions { (OptionalArg::Present(_), OptionalArg::Missing, OptionalArg::Present(_)) => { Err(vm.new_type_error(STRING_WITHOUT_ENCODING.to_owned())) } - } + }?; + Ok((inner, zero_count)) } } @@ -302,7 +315,7 @@ impl ByteInnerSub { } } - fn into_vec(self, vm: &VirtualMachine) -> PyResult> { + pub(crate) fn into_vec(self, vm: &VirtualMachine) -> PyResult> { Ok(match self { Self::Buffer(buffer) => buffer.elements, Self::Byte(int) => vec![int.as_bigint().byte_or(vm)?], @@ -321,6 +334,15 @@ pub struct ByteInnerFindOptions { } impl ByteInnerFindOptions { + pub(crate) fn parts(&self, len: usize) -> (&PyObject, core::ops::Range) { + let range = anystr::adjust_indices( + self.start.as_ref().map(AsRef::as_ref), + self.end.as_ref().map(AsRef::as_ref), + len, + ); + (&self.sub, range) + } + pub fn get_value( self, len: usize, @@ -347,7 +369,7 @@ pub struct ByteInnerPaddingOptions { } impl ByteInnerPaddingOptions { - fn get_value(self, fn_name: &str, vm: &VirtualMachine) -> PyResult<(isize, u8)> { + pub(crate) fn get_value(self, fn_name: &str, vm: &VirtualMachine) -> PyResult<(isize, u8)> { let fillchar = { let v = self.fillchar; try_as_bytes(v.clone(), |bytes| bytes.iter().copied().exactly_one().ok()) @@ -368,9 +390,9 @@ impl ByteInnerPaddingOptions { #[derive(FromArgs)] pub struct ByteInnerTranslateOptions { #[pyarg(positional)] - table: Option, + pub(crate) table: Option, #[pyarg(any, default = b"")] - delete: PyObjectRef, + pub(crate) delete: PyObjectRef, } impl ByteInnerTranslateOptions { @@ -708,6 +730,16 @@ impl PyBytesInner { vm: &VirtualMachine, ) -> PyResult> { let (width, fillchar) = options.get_value("center", vm)?; + self.pad_with(width, fillchar, pad, vm) + } + + pub(crate) fn pad_with( + &self, + width: isize, + fillchar: u8, + pad: PadFn, + vm: &VirtualMachine, + ) -> PyResult> { let len = self.len(); if len as isize >= width { return Ok(Vec::from(&self.elements[..])); @@ -783,7 +815,10 @@ impl PyBytesInner { vm: &VirtualMachine, ) -> PyResult> { let (table, delete) = options.get_value(vm)?; + Ok(self.translate_with(&table, &delete)) + } + pub(crate) fn translate_with(&self, table: &[u8], delete: &[u8]) -> Vec { let mut res = if delete.is_empty() { Vec::with_capacity(self.elements.len()) } else { @@ -796,7 +831,7 @@ impl PyBytesInner { } } - Ok(res) + res } pub fn strip(&self, chars: Option) -> Vec { diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 2c9d6075106..c7e69371006 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -6,10 +6,10 @@ use crate::{ AsObject, Py, PyExact, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, PyStackRef, TryFromObject, VirtualMachine, builtins::{ - PyBaseException, PyBaseExceptionRef, PyBaseObject, PyBytes, PyCode, PyCoroutine, PyDict, - PyDictRef, PyFloat, PyFrozenSet, PyGenerator, PyInt, PyInterpolation, PyList, PyModule, - PyProperty, PySet, PySlice, PyStr, PyStrInterned, PyTemplate, PyTraceback, PyType, - PyUtf8Str, + PyAnyDictRef, PyBaseException, PyBaseExceptionRef, PyBaseObject, PyBytes, PyCode, + PyCoroutine, PyDict, PyDictRef, PyFloat, PyFrozenSet, PyGenerator, PyInt, PyInterpolation, + PyList, PyModule, PyProperty, PySet, PySlice, PyStr, PyStrInterned, PyTemplate, + PyTraceback, PyType, PyUtf8Str, builtin_func::PyNativeFunction, descriptor::{PyMemberDescriptor, PyMethodDescriptor}, frame::stack_analysis, @@ -108,7 +108,7 @@ pub fn current_thread_frame_materialize(vm: &VirtualMachine) -> Option Option { +pub fn current_globals() -> Option { let ptr = crate::vm::thread::get_current_frame(); if ptr.is_null() { return None; @@ -923,7 +923,7 @@ pub struct InterpreterFrame { // Borrowed pointers — owned by FrameObject or by PyFunction on caller's stack. pub(crate) code: *const Py, pub(crate) func_obj: *const PyObject, // nullable - pub(crate) globals: *const Py, + pub(crate) globals: *const PyAnyDictRef, pub(crate) builtins: *const PyObject, /// Unified storage for local variables and evaluation stack. @@ -983,7 +983,7 @@ impl InterpreterFrame { #[inline(always)] pub(crate) unsafe fn new( code: &Py, - globals: &Py, + globals: &PyAnyDictRef, builtins: &PyObject, func_obj: Option<&PyObject>, localsplus: LocalsPlus, @@ -1027,7 +1027,7 @@ impl InterpreterFrame { Some(obj) => obj as *const PyObject, None => core::ptr::null(), }, - globals: globals as *const Py, + globals: globals as *const PyAnyDictRef, builtins: builtins as *const PyObject, localsplus, locals, @@ -1060,7 +1060,7 @@ impl InterpreterFrame { #[inline(always)] pub(crate) unsafe fn new_on_datastack<'a>( code: &Py, - globals: &Py, + globals: &PyAnyDictRef, builtins: &PyObject, func_obj: Option<&PyObject>, locals: FrameLocals, @@ -1266,7 +1266,7 @@ impl InterpreterFrame { // Create a full FrameObject with its own InterpreterFrame copy. // The FrameObject owns references to the same objects (code, globals, etc). let code: PyRef = self.code().to_owned(); - let globals: PyDictRef = self.globals().to_owned(); + let globals = self.globals().clone(); let builtins: PyObjectRef = self.builtins().to_owned(); let func_obj: Option = self.func_obj().map(|o| o.to_owned()); @@ -1362,7 +1362,7 @@ impl InterpreterFrame { #[cold] fn materialize_slow_chain(&self, vm: &VirtualMachine) -> FrameObjectRef { let code: PyRef = self.code().to_owned(); - let globals: PyDictRef = self.globals().to_owned(); + let globals = self.globals().clone(); let builtins: PyObjectRef = self.builtins().to_owned(); let func_obj: Option = self.func_obj().map(|o| o.to_owned()); @@ -1419,7 +1419,7 @@ impl InterpreterFrame { /// Borrowed globals dict. #[inline(always)] - pub fn globals(&self) -> &Py { + pub fn globals(&self) -> &PyAnyDictRef { // SAFETY: established by `new` / `init_iframe_ptrs`. unsafe { &*self.globals } } @@ -1534,7 +1534,7 @@ pub struct FrameObject { // raw pointers. Wrapped in Option so Traverse::clear can release them, // allowing GC cycle collection to reclaim referenced objects. pub(crate) owned_code: Option>, - pub(crate) owned_globals: Option, + pub(crate) owned_globals: Option, pub(crate) owned_builtins: Option, pub(crate) owned_func_obj: Option, @@ -1817,7 +1817,7 @@ impl FrameObject { let locals = match scope.locals { Some(locals) => FrameLocals::with_locals(locals), None if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) => FrameLocals::lazy(), - None => FrameLocals::with_locals(ArgMapping::from_dict_exact(scope.globals.clone())), + None => FrameLocals::with_locals(ArgMapping::from_anydict_exact(scope.globals.clone())), }; // Pointers are initially set from owned fields' references but will be @@ -1853,7 +1853,7 @@ impl FrameObject { fn init_iframe_ptrs(self_: &Py) { let iframe = unsafe { self_.iframe_mut() }; iframe.code = &**self_.owned_code.as_ref().unwrap() as *const Py; - iframe.globals = &**self_.owned_globals.as_ref().unwrap() as *const Py; + iframe.globals = self_.owned_globals.as_ref().unwrap() as *const PyAnyDictRef; iframe.builtins = &**self_.owned_builtins.as_ref().unwrap() as *const PyObject; iframe.func_obj = match &self_.owned_func_obj { Some(obj) => &**obj as *const PyObject, @@ -2497,7 +2497,7 @@ impl Py { }, ) }; - let builtins_dict = if globals.class().is(vm.ctx.types.dict_type) { + let builtins_dict = if globals.as_object().class().is(vm.ctx.types.dict_type) { builtins .downcast_ref_if_exact::(vm) // SAFETY: downcast_ref_if_exact already verified exact type @@ -2739,7 +2739,7 @@ fn exec_iframe<'a>( }, ) }; - let builtins_dict = if globals.class().is(vm.ctx.types.dict_type) { + let builtins_dict = if globals.as_object().class().is(vm.ctx.types.dict_type) { builtins .downcast_ref_if_exact::(vm) .map(|d| unsafe { PyExact::ref_unchecked(d) }) @@ -2931,7 +2931,7 @@ pub(crate) struct ExecutingFrame<'a> { code: &'a Py, localsplus: &'a mut LocalsPlus, locals: &'a FrameLocals, - globals: &'a Py, + globals: &'a PyAnyDictRef, builtins: &'a PyObject, /// Cached downcast of builtins to PyDict for fast LOAD_GLOBAL. /// Only set when both globals and builtins are exact dict types (not @@ -4651,15 +4651,10 @@ impl ExecutingFrame<'_> { } Instruction::DeleteName { namei: idx } => { let name = self.code.names[idx.get(arg) as usize]; - let res = self.locals.mapping(vm).ass_subscript(name, None, vm); - - match res { - Ok(()) => {} - Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => { - return Err(name_error(name, vm)); - } - Err(e) => return Err(e), - } + self.locals + .mapping(vm) + .ass_subscript(name, None, vm) + .map_err(|_| name_error(name, vm))?; Ok(None) } Instruction::DeleteSubscr => self.execute_delete_subscript(vm), @@ -5254,7 +5249,7 @@ impl ExecutingFrame<'_> { match result { Ok(x) => self.push_value(x), Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => { - self.push_value(self.load_global_or_builtin(name, vm)?); + self.push_value(self.load_name_global_or_builtin(name, vm)?); } Err(e) => return Err(e), } @@ -8043,9 +8038,9 @@ impl ExecutingFrame<'_> { let cached_version = self.code.instructions.read_cache_u16(cache_base + 1); let cached_index = self.code.instructions.read_cache_u16(cache_base + 3); if cached_version != 0 - && let Some(x) = self - .globals - .get_item_by_index_and_keys_version(cached_version, cached_index) + && let Some(globals) = self.globals.as_mutable() + && let Some(x) = + globals.get_item_by_index_and_keys_version(cached_version, cached_index) { self.push_value(x); if (oparg & 1) != 0 { @@ -8069,7 +8064,8 @@ impl ExecutingFrame<'_> { let cached_index = self.code.instructions.read_cache_u16(cache_base + 3); if cached_globals_ver != 0 && cached_builtins_ver != 0 - && let Ok(current_globals_ver) = u16::try_from(self.globals.keys_version()) + && let Some(globals) = self.globals.as_mutable() + && let Ok(current_globals_ver) = u16::try_from(globals.keys_version()) && cached_globals_ver == current_globals_ver && let Some(builtins_dict) = self.builtins.downcast_ref_if_exact::(vm) && let Some(x) = builtins_dict @@ -8541,7 +8537,12 @@ impl ExecutingFrame<'_> { if let Some(builtins_dict) = self.builtins_dict { // Fast path: both globals and builtins are exact dicts // SAFETY: builtins_dict is only set when globals is also exact dict - let globals_exact = unsafe { PyExact::ref_unchecked(self.globals) }; + let globals = self + .globals + .as_object() + .downcast_ref::() + .expect("exact dict globals"); + let globals_exact = unsafe { PyExact::ref_unchecked(globals) }; globals_exact .get_chain_exact(builtins_dict, name, vm)? .ok_or_else(|| { @@ -8552,16 +8553,29 @@ impl ExecutingFrame<'_> { if let Some(value) = self.globals.get_item_opt(name, vm)? { return Ok(value); } - self.builtins.get_item(name, vm).map_err(|e| { - if e.fast_isinstance(vm.ctx.exceptions.key_error) { - vm.new_name_error(format!("name '{name}' is not defined"), name.to_owned()) - } else { - e - } - }) + self.load_builtin(name, vm) } } + fn load_name_global_or_builtin(&self, name: &Py, vm: &VirtualMachine) -> PyResult { + // LOAD_NAME reads the stored global entries even for dict subclasses; + // LOAD_GLOBAL instead honors their __getitem__ implementation. + if let Some(value) = self.globals.inner_getitem_opt(name, vm)? { + return Ok(value); + } + self.load_builtin(name, vm) + } + + fn load_builtin(&self, name: &Py, vm: &VirtualMachine) -> PyResult { + self.builtins.get_item(name, vm).map_err(|e| { + if e.fast_isinstance(vm.ctx.exceptions.key_error) { + vm.new_name_error(format!("name '{name}' is not defined"), name.to_owned()) + } else { + e + } + }) + } + #[cfg_attr(feature = "flame-it", flame("FrameObject"))] fn import(&mut self, vm: &VirtualMachine, module_name: Option<&Py>) -> PyResult<()> { let module_name = module_name.unwrap_or(vm.ctx.empty_str); @@ -11786,7 +11800,7 @@ impl ExecutingFrame<'_> { let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) { FrameLocals::lazy() } else { - FrameLocals::with_locals(crate::function::ArgMapping::from_dict_exact( + FrameLocals::with_locals(crate::function::ArgMapping::from_anydict_exact( func.globals.clone(), )) }; @@ -11852,7 +11866,7 @@ impl ExecutingFrame<'_> { let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) { FrameLocals::lazy() } else { - FrameLocals::with_locals(crate::function::ArgMapping::from_dict_exact( + FrameLocals::with_locals(crate::function::ArgMapping::from_anydict_exact( func.globals.clone(), )) }; @@ -11938,8 +11952,11 @@ impl ExecutingFrame<'_> { ) { return; } + let Some(globals) = self.globals.as_mutable() else { + return; + }; let name = self.code.names[(oparg >> 1) as usize]; - let Ok(globals_version @ 1..) = u16::try_from(self.globals.assign_keys_version(vm)) else { + let Ok(globals_version @ 1..) = u16::try_from(globals.assign_keys_version(vm)) else { unsafe { self.code.instructions.write_adaptive_counter( cache_base, @@ -11951,7 +11968,7 @@ impl ExecutingFrame<'_> { return; }; - if let Ok(Some(globals_hint)) = self.globals.hint_for_key(name, vm) { + if let Ok(Some(globals_hint)) = globals.hint_for_key(name, vm) { unsafe { self.code .instructions diff --git a/crates/vm/src/function/protocol.rs b/crates/vm/src/function/protocol.rs index 5883a2a4691..f2f6ad370d0 100644 --- a/crates/vm/src/function/protocol.rs +++ b/crates/vm/src/function/protocol.rs @@ -1,7 +1,7 @@ use super::IntoFuncArgs; use crate::{ AsObject, PyObject, PyObjectRef, PyPayload, PyResult, TryFromObject, VirtualMachine, - builtins::{PyDictRef, iter::PySequenceIterator}, + builtins::{PyAnyDictRef, PyDictRef, iter::PySequenceIterator}, convert::ToPyObject, object::{Traverse, TraverseFn}, protocol::{PyIter, PyIterIter, PyMapping}, @@ -155,6 +155,12 @@ impl ArgMapping { Self { obj: dict.into() } } + #[inline(always)] + #[must_use] + pub fn from_anydict_exact(dict: PyAnyDictRef) -> Self { + Self { obj: dict.into() } + } + #[inline(always)] #[must_use] pub fn obj(&self) -> &PyObject { diff --git a/crates/vm/src/protocol/callable.rs b/crates/vm/src/protocol/callable.rs index 49172a15222..2dcec1838c0 100644 --- a/crates/vm/src/protocol/callable.rs +++ b/crates/vm/src/protocol/callable.rs @@ -2,7 +2,7 @@ use crate::{ builtins::{PyBoundMethod, PyFunction}, function::{FuncArgs, IntoFuncArgs}, types::{GenericMethod, VectorCallFunc}, - {PyObject, PyObjectRef, PyResult, VirtualMachine}, + {AsObject, PyObject, PyObjectRef, PyResult, VirtualMachine}, }; impl PyObject { @@ -78,6 +78,15 @@ impl<'a> PyCallable<'a> { pub fn invoke(&self, args: impl IntoFuncArgs, vm: &VirtualMachine) -> PyResult { let args = args.into_args(vm); + // The frozendict identity optimization belongs to the ordinary call path, not + // to __new__ or an explicit type.__call__. CALL_FUNCTION_EX reaches this path. + if self.obj.is(vm.ctx.types.frozendict_type) + && args.kwargs.is_empty() + && self.vectorcall.is_some() + { + let nargs = args.args.len(); + return self.invoke_vectorcall(args.args, nargs, None, vm); + } if !vm.use_tracing.get() { return (self.call)(self.obj, args, vm); } diff --git a/crates/vm/src/protocol/mapping.rs b/crates/vm/src/protocol/mapping.rs index e791757901e..997df878ff0 100644 --- a/crates/vm/src/protocol/mapping.rs +++ b/crates/vm/src/protocol/mapping.rs @@ -184,7 +184,7 @@ impl PyMapping<'_> { pub fn keys(self, vm: &VirtualMachine) -> PyResult { if let Some(dict) = self.obj.downcast_ref_if_exact::(vm) { - PyDictKeys::new(dict.to_owned()).to_pyresult(vm) + PyDictKeys::new(dict.to_owned().into()).to_pyresult(vm) } else { self.method_output_as_list(identifier!(vm, keys), vm) } @@ -192,7 +192,7 @@ impl PyMapping<'_> { pub fn values(self, vm: &VirtualMachine) -> PyResult { if let Some(dict) = self.obj.downcast_ref_if_exact::(vm) { - PyDictValues::new(dict.to_owned()).to_pyresult(vm) + PyDictValues::new(dict.to_owned().into()).to_pyresult(vm) } else { self.method_output_as_list(identifier!(vm, values), vm) } @@ -200,7 +200,7 @@ impl PyMapping<'_> { pub fn items(self, vm: &VirtualMachine) -> PyResult { if let Some(dict) = self.obj.downcast_ref_if_exact::(vm) { - PyDictItems::new(dict.to_owned()).to_pyresult(vm) + PyDictItems::new(dict.to_owned().into()).to_pyresult(vm) } else { self.method_output_as_list(identifier!(vm, items), vm) } diff --git a/crates/vm/src/scope.rs b/crates/vm/src/scope.rs index 772d5c176c8..ed67b6830b7 100644 --- a/crates/vm/src/scope.rs +++ b/crates/vm/src/scope.rs @@ -1,10 +1,14 @@ -use crate::{VirtualMachine, builtins::PyDictRef, function::ArgMapping}; +use crate::{ + VirtualMachine, + builtins::{PyAnyDictRef, PyDictRef}, + function::ArgMapping, +}; use alloc::fmt; #[derive(Clone)] pub struct Scope { pub locals: Option, - pub globals: PyDictRef, + pub globals: PyAnyDictRef, } impl fmt::Debug for Scope { @@ -17,8 +21,11 @@ impl fmt::Debug for Scope { impl Scope { #[inline] #[must_use] - pub const fn new(locals: Option, globals: PyDictRef) -> Self { - Self { locals, globals } + pub fn new(locals: Option, globals: impl Into) -> Self { + Self { + locals, + globals: globals.into(), + } } pub fn with_builtins( diff --git a/crates/vm/src/stdlib/_io.rs b/crates/vm/src/stdlib/_io.rs index 5bcd8f64d29..d20cb367999 100644 --- a/crates/vm/src/stdlib/_io.rs +++ b/crates/vm/src/stdlib/_io.rs @@ -762,10 +762,13 @@ mod _io { ))); } let n = n as usize; - let mut bytes = b.borrow_buf_mut(); - bytes.truncate(n); - // FIXME: try to use Arc::unwrap on the bytearray to get at the inner buffer - bytes.clone().to_pyobject(vm) + let mut inner = b.try_resizable(vm)?; + let bytes = &mut inner.elements; + bytes + .try_reserve_exact(n.saturating_sub(bytes.len())) + .map_err(|_| vm.no_memory_error())?; + bytes.resize(n, 0); + core::mem::take(bytes).to_pyobject(vm) } }) } else { diff --git a/crates/vm/src/stdlib/_testinternalcapi.rs b/crates/vm/src/stdlib/_testinternalcapi.rs index caafcf5ab57..24f790412ab 100644 --- a/crates/vm/src/stdlib/_testinternalcapi.rs +++ b/crates/vm/src/stdlib/_testinternalcapi.rs @@ -185,8 +185,8 @@ mod _testinternalcapi { &mut counts, args.globalnames.into_option(), args.attrnames.into_option(), - globalsns.as_deref(), - builtinsns.as_deref(), + globalsns.as_ref(), + builtinsns.as_ref(), vm, )?; counts_to_dict(&counts, vm) @@ -214,8 +214,8 @@ mod _testinternalcapi { &mut counts, args.globalnames.into_option(), None, - globalsns.as_deref(), - builtinsns.as_deref(), + globalsns.as_ref(), + builtinsns.as_ref(), vm, )?; if builtinsns.is_some() { @@ -1125,11 +1125,15 @@ fn expandtabs(input: &str, tab_size: usize) -> String { expanded } -type CodeNamespaces = (PyRef, Option>, Option>); +type CodeNamespaces = ( + PyRef, + Option, + Option, +); fn code_or_function(obj: &PyObject, vm: &VirtualMachine) -> PyResult { if let Ok(func) = obj.to_owned().downcast::() { - let builtins = func.builtins.clone().downcast::().ok(); + let builtins = crate::builtins::PyAnyDictRef::from_object(&func.builtins); return Ok(( (*func.code).to_owned(), Some(func.globals.clone()), @@ -1142,19 +1146,22 @@ fn code_or_function(obj: &PyObject, vm: &VirtualMachine) -> PyResult>>( override_ns: OptionalArg, - default: Option>, + default: Option, name: &str, vm: &VirtualMachine, -) -> PyResult>> { +) -> PyResult> { match override_ns.into_option() { - Some(obj) => obj.downcast::().map(Some).map_err(|obj| { - vm.new_type_error(format!( - "expected a dict for \"{name}\", got {}", - obj.class().name() - )) - }), + Some(obj) => obj + .downcast::() + .map(|dict| Some(dict.into())) + .map_err(|obj| { + vm.new_type_error(format!( + "expected a dict for \"{name}\", got {}", + obj.class().name() + )) + }), None => Ok(default), } } @@ -1280,8 +1287,8 @@ fn set_unbound_var_counts( counts: &mut VarCounts, globalnames: Option, attrnames: Option, - globalsns: Option<&Py>, - builtinsns: Option<&Py>, + globalsns: Option<&crate::builtins::PyAnyDictRef>, + builtinsns: Option<&crate::builtins::PyAnyDictRef>, vm: &VirtualMachine, ) -> PyResult<()> { let globalnames = optional_set(globalnames, "globalnames", vm)?; @@ -1317,8 +1324,8 @@ fn identify_unbound_names( code: &Py, globalnames: Option>, attrnames: Option>, - globalsns: Option<&Py>, - builtinsns: Option<&Py>, + globalsns: Option<&crate::builtins::PyAnyDictRef>, + builtinsns: Option<&crate::builtins::PyAnyDictRef>, vm: &VirtualMachine, ) -> PyResult<(UnboundCounts, i32)> { let mut seen_globals = HashSet::new(); diff --git a/crates/vm/src/stdlib/builtins.rs b/crates/vm/src/stdlib/builtins.rs index 655f30f87d8..ab07eb634d3 100644 --- a/crates/vm/src/stdlib/builtins.rs +++ b/crates/vm/src/stdlib/builtins.rs @@ -11,7 +11,7 @@ mod builtins { AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine, builtins::{ - PyByteArray, PyBytes, PyDictRef, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType, + PyAnyDictRef, PyByteArray, PyBytes, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType, PyUtf8StrRef, enumerate::PyReverseSequenceIterator, function::{PyCell, PyCellRef, PyFunction}, @@ -436,21 +436,23 @@ mod builtins { vm: &VirtualMachine, func_name: &'static str, ) -> PyResult<()> { - if !globals.fast_isinstance(vm.ctx.types.dict_type) { + if !globals.fast_isinstance(vm.ctx.types.dict_type) + && !globals.fast_isinstance(vm.ctx.types.frozendict_type) + { return Err(match func_name { "eval" => { let is_mapping = globals.mapping_unchecked().check(); vm.new_type_error(if is_mapping { - "globals must be a real dict; try eval(expr, {}, mapping)" + "globals must be a real dict or a frozendict; try eval(expr, {}, mapping)" } else { - "globals must be a dict" + "globals must be a dict or a frozendict" }) } "exec" => vm.new_type_error(format!( - "exec() globals must be a dict, not {}", + "exec() globals must be a dict or a frozendict, not {}", globals.class().name() )), - _ => vm.new_type_error("globals must be a dict"), + _ => vm.new_type_error("globals must be a dict or a frozendict"), }); } Ok(()) @@ -460,15 +462,11 @@ mod builtins { Some(globals) => { validate_globals_dict(&globals, vm, func_name)?; - let globals = PyDictRef::try_from_object(vm, globals)?; - if !globals.contains_key(identifier!(vm, __builtins__), vm) { - let builtins_dict = vm.builtins.dict().into(); - globals.set_item(identifier!(vm, __builtins__), builtins_dict, vm)?; - } + let globals = PyAnyDictRef::try_from_object(vm, globals)?; ( globals.clone(), self.locals - .unwrap_or_else(|| ArgMapping::from_dict_exact(globals.clone())), + .unwrap_or_else(|| ArgMapping::from_anydict_exact(globals.clone())), ) } None => ( @@ -481,7 +479,27 @@ mod builtins { ), }; - let scope = crate::scope::Scope::with_builtins(Some(locals), globals, vm); + let builtins = globals.inner_getitem_opt(identifier!(vm, __builtins__), vm)?; + if builtins.is_none() { + if globals.is_frozen() { + let builtins = crate::frame::current_builtins() + .unwrap_or_else(|| vm.builtins.dict().into()); + globals + .as_object() + .set_item(identifier!(vm, __builtins__), builtins, vm) + .map_err(|_| { + vm.new_type_error("cannot assign __builtins__ to frozendict globals") + })?; + } else { + globals.set_item( + identifier!(vm, __builtins__), + vm.builtins.dict().into(), + vm, + )?; + } + } + + let scope = crate::scope::Scope::new(Some(locals), globals); Ok(scope) } } @@ -650,6 +668,7 @@ mod builtins { func: &str, closure: Option>>, ) -> PyResult { + let from_string = matches!(&source, Either::B(_)); // Determine code object: let code_obj = match source { #[cfg(feature = "rustpython-compiler")] @@ -675,6 +694,20 @@ mod builtins { ))); } + // The source-string path repeats the stored-builtins check through the + // dictionary API, unlike direct evaluation of an existing code object. + if from_string + && scope.globals.is_frozen() + && scope + .globals + .inner_getitem_opt(identifier!(vm, __builtins__), vm)? + .is_none() + { + scope + .globals + .set_item(identifier!(vm, __builtins__), vm.builtins.dict().into(), vm)?; + } + // Run the code: vm.run_code_obj_with_closure(code_obj, scope, closure) } @@ -721,7 +754,7 @@ mod builtins { } #[pyfunction] - fn globals(vm: &VirtualMachine) -> PyDictRef { + fn globals(vm: &VirtualMachine) -> PyAnyDictRef { vm.current_globals() } @@ -1490,6 +1523,7 @@ pub fn init_module(vm: &VirtualMachine, module: &Py) { "dict" => ctx.types.dict_type.to_owned(), "enumerate" => ctx.types.enumerate_type.to_owned(), "float" => ctx.types.float_type.to_owned(), + "frozendict" => ctx.types.frozendict_type.to_owned(), "frozenset" => ctx.types.frozenset_type.to_owned(), "filter" => ctx.types.filter_type.to_owned(), "int" => ctx.types.int_type.to_owned(), diff --git a/crates/vm/src/types/zoo.rs b/crates/vm/src/types/zoo.rs index fce1a1e9eb6..5291e677add 100644 --- a/crates/vm/src/types/zoo.rs +++ b/crates/vm/src/types/zoo.rs @@ -40,6 +40,7 @@ pub struct TypeZoo { pub float_type: &'static Py, pub frame_type: &'static Py, pub frame_locals_proxy_type: &'static Py, + pub frozendict_type: &'static Py, pub frozenset_type: &'static Py, pub generator_type: &'static Py, pub int_type: &'static Py, @@ -139,6 +140,7 @@ impl TypeZoo { dict_type: dict::PyDict::init_builtin_type(), enumerate_type: enumerate::PyEnumerate::init_builtin_type(), float_type: float::PyFloat::init_builtin_type(), + frozendict_type: dict::PyFrozenDict::init_builtin_type(), frozenset_type: set::PyFrozenSet::init_builtin_type(), filter_type: filter::PyFilter::init_builtin_type(), list_type: list::PyList::init_builtin_type(), diff --git a/crates/vm/src/vm/crossinterp.rs b/crates/vm/src/vm/crossinterp.rs index fbab68f153d..a72d34ea92a 100644 --- a/crates/vm/src/vm/crossinterp.rs +++ b/crates/vm/src/vm/crossinterp.rs @@ -340,7 +340,7 @@ impl SharedValue { // Stateless functions have no globals, so `__main__` is used, // just like for builtins such as exec(). let globals = vm.main_namespace()?; - Ok(PyFunction::new(code, globals, vm)?.into_pyobject(vm)) + Ok(PyFunction::new(code, globals.into(), vm)?.into_pyobject(vm)) } Self::Pickled(data) => pickle_loads(&data, vm), } @@ -763,7 +763,12 @@ pub(crate) fn walk_instructions( /// `_PyFunction_VerifyStateless`. fn verify_stateless_function(func: &Py, vm: &VirtualMachine) -> PyResult<()> { - // `__globals__` is a dict by construction, so only the builtins are checked. + let globals = func.globals.as_mutable().ok_or_else(|| { + vm.new_type_error(format!( + "unsupported globals {}", + render_repr(func.globals.as_object(), vm) + )) + })?; let builtins = func.builtins.downcast_ref::().ok_or_else(|| { vm.new_type_error(format!( "unsupported builtins {}", @@ -786,7 +791,7 @@ fn verify_stateless_function(func: &Py, vm: &VirtualMachine) -> PyRe { return Err(vm.new_value_error("closures not supported")); } - verify_stateless(&func.code, Some((&func.globals, builtins)), vm) + verify_stateless(&func.code, Some((globals, builtins)), vm) } /// `verify_script`: a script takes no arguments and returns only None. diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 304991f9504..b8928227761 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -24,7 +24,7 @@ mod vm_ops; use crate::{ AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, builtins::{ - self, PyBaseExceptionRef, PyBaseObject, PyDict, PyDictRef, PyFrozenSet, PyInt, PyList, + self, PyAnyDictRef, PyBaseExceptionRef, PyBaseObject, PyDict, PyFrozenSet, PyInt, PyList, PyModule, PySet, PyStr, PyStrInterned, PyStrRef, PyTypeRef, PyUtf8Str, PyUtf8StrInterned, PyWeak, code::PyCode, @@ -3113,7 +3113,7 @@ impl VirtualMachine { .locals(self) } - pub fn current_globals(&self) -> PyDictRef { + pub fn current_globals(&self) -> PyAnyDictRef { let ptr = crate::vm::thread::get_current_frame(); if !ptr.is_null() { return unsafe { (*ptr).globals().to_owned() }; @@ -3226,7 +3226,7 @@ impl VirtualMachine { // Locals fallback: use the heavy frame if available, otherwise // use globals as locals (light frame locals are on the data stack). let locals_mapping = self.current_frame().map_or_else( - || ArgMapping::from_dict_exact(globals.clone()), + || ArgMapping::from_anydict_exact(globals.clone()), |f| f.iframe().locals.clone_mapping(self), ); (Some(locals_mapping), Some(globals)) diff --git a/crates/vm/src/vm/vm_new.rs b/crates/vm/src/vm/vm_new.rs index 1a1696b7106..7c5d43d1643 100644 --- a/crates/vm/src/vm/vm_new.rs +++ b/crates/vm/src/vm/vm_new.rs @@ -83,15 +83,16 @@ impl VirtualMachine { } pub fn new_scope_with_main(&self) -> PyResult { - let scope = self.new_scope_with_builtins(); - let main_module = self.new_module("__main__", scope.globals.clone(), None); + let globals = self.ctx.new_dict(); + let scope = Scope::with_builtins(None, globals.clone(), self); + let main_module = self.new_module("__main__", globals.clone(), None); self.sys_module.get_attr("modules", self)?.set_item( "__main__", main_module.into(), self, )?; - self.set_main_builtin_importer(&scope.globals)?; + self.set_main_builtin_importer(&globals)?; Ok(scope) } diff --git a/crates/vm/src/warn.rs b/crates/vm/src/warn.rs index 3d8e7cb6fb1..52e440a9d64 100644 --- a/crates/vm/src/warn.rs +++ b/crates/vm/src/warn.rs @@ -762,12 +762,19 @@ fn setup_context( d.downcast::() .map_err(|_| vm.new_type_error("sys.__dict__ is not a dictionary")) })?; - (globals, vm.ctx.intern_str(""), 0) + (globals.into(), vm.ctx.intern_str(""), 0) }; - let registry = match globals.get_item("__warningregistry__", vm) { - Ok(r) => r, - Err(_) => { + let get_global = |name: &str| { + if globals.is_frozen() { + globals.inner_getitem_opt(name, vm) + } else { + Ok(globals.get_item(name, vm).ok()) + } + }; + let registry = match get_global("__warningregistry__")? { + Some(r) => r, + None => { let r = vm.ctx.new_dict(); globals.set_item("__warningregistry__", r.clone().into(), vm)?; r.into() @@ -775,8 +782,6 @@ fn setup_context( }; // Setup module. - let module = globals - .get_item("__name__", vm) - .unwrap_or_else(|_| vm.new_pyobj("")); + let module = get_global("__name__")?.unwrap_or_else(|| vm.new_pyobj("")); Ok((filename.to_owned(), lineno, Some(module), registry)) } diff --git a/extra_tests/snippets/builtin_bytearray.py b/extra_tests/snippets/builtin_bytearray.py index d897ec94782..6b397b12fcf 100644 --- a/extra_tests/snippets/builtin_bytearray.py +++ b/extra_tests/snippets/builtin_bytearray.py @@ -913,3 +913,643 @@ class B(bytearray): empty = bytearray() empty *= sys.maxsize assert empty == bytearray() + + +def test_bytearray_exports(): + """Resizing paths must respect outstanding bytearray exports.""" + + import io + + from testutils import assert_raises + + for args in ((), (b"abc",), (b"x",), (b"",), (object(),)): + value = bytearray(b"abc") + with memoryview(value) as view: + with assert_raises(BufferError): + value.__init__(*args) + assert value == b"abc" + assert view.tobytes() == b"abc" + value.__init__(b"released") + assert value == b"released" + + value = bytearray(b"abc") + with assert_raises(TypeError): + value.__init__(object()) + assert value == b"" + value = bytearray(b"abc") + value.__init__(value) + assert value == b"" + + class ExportDuringInit: + def __init__(self, count, payload=b""): + self.count = count + self.payload = payload + self.calls = 0 + + def __index__(self): + self.calls += 1 + value.extend(self.payload) + self.view = memoryview(value) + return self.count + + value = bytearray(b"abc") + source = ExportDuringInit(2) + with assert_raises(BufferError): + value.__init__(source) + assert value == b"" + assert source.view.tobytes() == b"" + assert source.calls == 1 + source.view.release() + value.extend(b"released") + assert value == b"released" + + value = bytearray(b"abc") + source = ExportDuringInit(0) + value.__init__(source) + assert value == b"" + assert source.view.tobytes() == b"" + source.view.release() + + value = bytearray(b"abc") + source = ExportDuringInit(0, b"xy") + value.__init__(source) + assert value == b"xy" + assert source.view.tobytes() == b"xy" + assert source.calls == 1 + source.view.release() + + class FailingIndex: + def __index__(self): + assert value == b"" + raise ValueError("conversion failed") + + value = bytearray(b"abc") + with assert_raises(ValueError): + value.__init__(FailingIndex()) + assert value == b"" + + for key in (0, slice(0, 1)): + value = bytearray(b"abc") + with memoryview(value) as view: + with assert_raises(BufferError): + del value[key] + assert view.tobytes() == b"abc" + del value[key] + assert value == b"bc" + + class Reader(io.RawIOBase): + def __init__(self, count, retain, payload=b"abcd"): + self.count = count + self.retain = retain + self.payload = payload + + def readinto(self, buffer): + buffer[:] = self.payload + self.buffer = buffer + if self.retain: + self.view = memoryview(buffer) + return self.count + + for count in (0, 2, 4): + reader = Reader(count, False) + assert reader.read(4) == b"abcd"[:count] + assert reader.buffer == b"" + + reader = Reader(count, True) + with assert_raises(BufferError): + reader.read(4) + assert reader.view.tobytes() == b"abcd" + reader.view.release() + reader.buffer.clear() + + for count in (-1, 5): + reader = Reader(count, True) + with assert_raises(ValueError): + reader.read(4) + assert reader.view.tobytes() == b"abcd" + reader.view.release() + + reader = Reader(None, True) + assert reader.read(4) is None + assert reader.view.tobytes() == b"abcd" + reader.view.release() + + # The callback can shrink the temporary buffer, so slicing by the original + # requested size without checking the current buffer is not safe. + reader = Reader(2, False, b"xy") + assert reader.read(4) == b"xy" + assert reader.buffer == b"" + + reader = Reader(4, False, b"xy") + result = reader.read(4) + assert len(result) == 4 + assert result[:2] == b"xy" + # The contents of the newly grown region are unspecified by CPython. + assert reader.buffer == b"" + + +test_bytearray_exports() + + +def test_bytearray_threading(): + """Acquiring a bytearray export must be atomic with respect to resizing.""" + + import threading + + def exercise_exports(): + payload = b"x" * 65536 + value = bytearray(payload) + barrier = threading.Barrier(3, timeout=5) + errors = [] + + def worker(operation): + try: + for _ in range(256): + barrier.wait() + operation() + barrier.wait() + except BaseException as exc: + errors.append(exc) + barrier.abort() + + def resize(): + try: + value.clear() + value.extend(payload) + except BufferError: + pass + + def decode(): + result = value.decode("ascii") + assert result in ("", "x" * len(payload)) + + def export(): + with memoryview(value) as view: + size = len(view) + assert size in (0, len(payload)) + # The acquired export keeps its original length and contents even + # if the resizing thread runs between these two reads. + assert view.tobytes() == payload[:size] + assert len(view) == size + + threads = [ + threading.Thread(target=worker, args=(operation,), daemon=True) + for operation in (resize, decode, export) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + assert not any(thread.is_alive() for thread in threads), "worker did not finish" + if errors: + raise errors[0] + + exercise_exports() + + +test_bytearray_threading() + + +def test_bytearray_init_exports(): + """Initializer storage and integer-deletion ordering with buffer exports.""" + + from testutils import assert_raises + + class Zero: + def __init__(self, value): + self.value = value + self.observations = [] + + def __index__(self): + self.observations.append((len(self.value), self.value.__alloc__())) + return 0 + + # Final CPython rejects even an empty exported initializer before conversion. + for kind in ("missing", "integer", "callback"): + value = bytearray() + source = Zero(value) + with memoryview(value) as view: + with assert_raises(BufferError): + if kind == "missing": + value.__init__() + elif kind == "integer": + value.__init__(0) + else: + value.__init__(source) + assert value == b"" + assert view.tobytes() == b"" + assert source.observations == [] + + # Clearing an initializer must release its old storage before callbacks and + # before the successful zero-count early return. + for kind in ("missing", "integer", "callback"): + value = bytearray(b"x" * 64) + source = Zero(value) + assert value.__alloc__() > 0 + if kind == "missing": + value.__init__() + elif kind == "integer": + value.__init__(0) + else: + value.__init__(source) + assert source.observations == [(0, 0)] + assert value == b"" + assert value.__alloc__() == 0 + + # An invalid integer index is rejected before checking whether resizing is + # blocked, and neither error path may alter the exported bytes. + for initial, indices in ((b"abc", (3, -4)), (b"", (0, -1))): + for index in indices: + value = bytearray(initial) + with memoryview(value) as view: + try: + del value[index] + except IndexError as exc: + assert str(exc) == "bytearray index out of range" + else: + raise AssertionError("invalid index was accepted") + assert value == initial + assert view.tobytes() == initial + + for index in (0, -1): + value = bytearray(b"abc") + with memoryview(value) as view: + with assert_raises(BufferError): + del value[index] + assert view.tobytes() == b"abc" + del value[index] + assert value == (b"bc" if index == 0 else b"ab") + + +test_bytearray_init_exports() + + +def test_bytearray_concat_exports(): + import operator + + from testutils import assert_raises + + for concat in (operator.add, operator.concat, bytearray.__add__): + target = bytearray(b"a") + events = [] + + class ReentrantBuffer: + def __buffer__(self, flags): + events.append("acquire") + with assert_raises(BufferError): + target.__init__(object()) + with assert_raises(BufferError): + target.append(120) + target[0] = ord("A") + return memoryview(b"b") + + def __release_buffer__(self, view): + events.append("release") + # The left operand's export must already have been released. + target.append(99) + + assert concat(target, ReentrantBuffer()) == bytearray(b"Ab") + assert events == ["acquire", "release"] + assert target == bytearray(b"Ac") + target.append(100) + + class FailingBuffer: + def __buffer__(self, flags): + raise ValueError("buffer acquisition failed") + + with assert_raises(TypeError) as error: + concat(target, FailingBuffer()) + assert str(error.exception) == "can't concat FailingBuffer to bytearray" + target.append(101) + + events = [] + + class Left(bytearray): + def __buffer__(self, flags): + events.append("left acquire") + return memoryview(b"L") + + def __release_buffer__(self, view): + events.append("left release") + + class Right: + def __buffer__(self, flags): + events.append("right acquire") + return memoryview(b"R") + + def __release_buffer__(self, view): + events.append("right release") + + assert concat(Left(b"ignored"), Right()) == bytearray(b"LR") + assert events == [ + "left acquire", + "right acquire", + "left release", + "right release", + ] + + events = [] + with assert_raises(TypeError) as error: + concat(Left(), FailingBuffer()) + assert str(error.exception) == "can't concat FailingBuffer to Left" + assert events == ["left acquire", "left release"] + + class FailingLeft(bytearray): + def __buffer__(self, flags): + events.append("left acquire") + raise ValueError("left acquisition failed") + + events = [] + with assert_raises(TypeError) as error: + concat(FailingLeft(), Right()) + assert str(error.exception) == "can't concat Right to FailingLeft" + assert events == ["left acquire"] + + +test_bytearray_concat_exports() + + +def test_bytearray_callback_exports(): + import operator + + from testutils import assert_raises + + def check_comparison(): + for compare, expected in ( + (operator.eq, False), + (operator.ne, True), + (operator.lt, True), + (operator.le, True), + (operator.gt, False), + (operator.ge, False), + ): + target = bytearray(b"a") + events = [] + + class Other: + def __buffer__(self, flags): + events.append("acquire") + with assert_raises(BufferError): + target.__init__(object()) + target[0] = ord("A") + return memoryview(b"B") + + def __release_buffer__(self, view): + events.append("release") + target[0] = ord("Z") + target.append(99) + + assert compare(target, Other()) is expected + assert target == b"Zc" + assert events == ["acquire", "release"] + + events = [] + + class Left(bytearray): + def __buffer__(self, flags): + events.append("acquire") + return memoryview(b"L") + + def __release_buffer__(self, view): + events.append("release") + + left = Left(b"ignored") + assert left.__eq__(left) is True + assert events == ["acquire", "acquire", "release", "release"] + events.clear() + assert left.__eq__(object()) is NotImplemented + assert events == [] + + class Failing: + def __buffer__(self, flags): + raise ValueError("no buffer") + + assert left.__eq__(Failing()) is NotImplemented + assert events == ["acquire", "release"] + left.append(0) + + def check_search(): + for method, expected in ( + ("find", 0), + ("index", 0), + ("rfind", 2), + ("rindex", 2), + ("count", 2), + ): + target = bytearray(b"xxx") + events = [] + + class Needle: + def __buffer__(self, flags): + events.append("acquire") + with assert_raises(BufferError): + target.__init__(object()) + target[:] = b"aba" + return memoryview(b"a") + + def __release_buffer__(self, view): + events.append("release") + with assert_raises(BufferError): + target.append(0) + target[:] = b"zzz" + + assert getattr(target, method)(Needle()) == expected + assert target == b"zzz" + assert events == ["acquire", "release"] + target.append(0) + + target = bytearray(b"xxx") + + class IndexNeedle: + def __index__(self): + with assert_raises(BufferError): + target.__init__(object()) + target[:] = b"aba" + return ord("a") + + assert getattr(target, method)(IndexNeedle()) == expected + + target = bytearray(b"ab") + + class Start: + def __index__(self): + target.extend(b"ab") + return 1 + + assert target.count(ord("a"), Start()) == 1 + assert target == b"abab" + + def check_affixes(): + for method in ("startswith", "endswith"): + target = bytearray(b"aaa") + events = [] + + class Affix: + def __init__(self, value, after): + self.value = value + self.after = after + + def __buffer__(self, flags): + events.append(self.value) + with assert_raises(BufferError): + target.__init__(object()) + target[0] = target[0] + return memoryview(self.value) + + def __release_buffer__(self, view): + events.append("release") + with assert_raises(BufferError): + target.append(0) + target[:] = self.after + + class Unvisited: + def __buffer__(self, flags): + raise AssertionError("tuple must short-circuit") + + assert getattr(target, method)( + (Affix(b"x", b"yyy"), Affix(b"y", b"zzz"), Unvisited()) + ) + assert target == b"zzz" + assert events == [b"x", "release", b"y", "release"] + target.append(0) + + target = bytearray(b"aaa") + events.clear() + assert getattr(target, method)(Affix(b"a", b"zzz")) + assert target == b"zzz" + assert events == [b"a", "release"] + + events.clear() + assert not getattr(target, method)(Affix(b"z", b"aaa"), 100) + assert events == [b"z", "release"] + with assert_raises(TypeError): + getattr(target, method)((Affix(b"x", b"aaa"), object())) + target.append(0) + + class BytesAffix(bytes): + def __buffer__(self, flags): + raise AssertionError("bytes subclasses use their stored bytes") + + assert getattr(bytearray(b"a"), method)(BytesAffix(b"a")) + + def check_translate(): + target = bytearray(b"x") + table = bytearray(range(256)) + events = [] + + class Table: + def __buffer__(self, flags): + events.append("table acquire") + target.__init__(b"a") + return memoryview(table) + + def __release_buffer__(self, view): + events.append("table release") + target.__init__(b"t") + + class Delete: + def __buffer__(self, flags): + events.append("delete acquire") + with assert_raises(BufferError): + table.append(0) + table[ord("a")] = ord("b") + target.__init__(b"aa") + return memoryview(b"") + + def __release_buffer__(self, view): + events.append("delete release") + target.__init__(b"d") + + assert target.translate(Table(), delete=Delete()) == b"bb" + assert target == b"d" + assert events == [ + "table acquire", + "delete acquire", + "table release", + "delete release", + ] + table.append(0) + table.pop() + events.clear() + + class BadDelete: + def __buffer__(self, flags): + events.append("delete failure") + with assert_raises(BufferError): + table.append(0) + raise ValueError("bad delete") + + with assert_raises(ValueError): + target.translate(Table(), delete=BadDelete()) + assert events == ["table acquire", "delete failure", "table release"] + table.append(0) + + def check_padding(): + for method in ("center", "ljust", "rjust"): + target = bytearray(b"x") + assert getattr(target, method)(3, target) == b"xxx" + + class Width: + def __index__(self): + target.__init__(b"y") + return 3 + + assert getattr(target, method)(Width(), target) == b"yyy" + + check_comparison() + check_search() + check_affixes() + check_translate() + check_padding() + + +test_bytearray_callback_exports() + + +def test_bytearray_missing_source_init(): + from testutils import assert_raises + + for initial in (b"", b"previous contents"): + value = bytearray(initial) + assert value.__init__() is None + assert value == b"" + assert value.__alloc__() == 0 + + for kwargs, message in ( + ({"encoding": "utf-8"}, "encoding without a string argument"), + ({"errors": "strict"}, "errors without a string argument"), + ( + {"encoding": "utf-8", "errors": "strict"}, + "encoding without a string argument", + ), + ): + value = bytearray(initial) + with assert_raises(TypeError) as caught: + value.__init__(**kwargs) + assert str(caught.exception) == message + assert value == b"" + assert value.__alloc__() == 0 + + for kwargs in ( + {}, + {"encoding": "utf-8"}, + {"errors": "strict"}, + {"encoding": "utf-8", "errors": "strict"}, + ): + value = bytearray(initial) + with memoryview(value) as view: + with assert_raises(BufferError): + value.__init__(**kwargs) + assert value == initial + assert view.tobytes() == initial + assert value.__init__() is None + assert value == b"" + + # Explicit None is a source and must still go through conversion. + value = bytearray(initial) + with assert_raises(TypeError): + value.__init__(None) + assert value == b"" + + +test_bytearray_missing_source_init() diff --git a/extra_tests/snippets/builtin_dict.py b/extra_tests/snippets/builtin_dict.py index e593b92fd87..220dd6b321a 100644 --- a/extra_tests/snippets/builtin_dict.py +++ b/extra_tests/snippets/builtin_dict.py @@ -26,6 +26,7 @@ next(dict_reversed) assert "dict" in dict().__doc__ + d = {"a": 123, "b": 456} assert 1 not in d.items() assert "a" not in d.items() @@ -873,3 +874,1356 @@ def __hash__(self): # entirely when popping from an empty dict). {1: 1}.pop(CountingHash()) assert CountingHash.calls == 1, CountingHash.calls + + +# PEP 814 regressions, validated against CPython 3.15.0. +import builtins + +if hasattr(builtins, "frozendict"): + assert ( + frozendict.__doc__ + == """frozendict() -> new empty immutable dictionary +frozendict(mapping) -> new immutable dictionary initialized from a mapping + object's (key, value) pairs +frozendict(iterable) -> new immutable dictionary initialized as if via: + d = {} + for k, v in iterable: + d[k] = v + d = frozendict(d) +frozendict(**kwargs) -> new immutable dictionary initialized with the name=value + pairs in the keyword argument list. For example: frozendict(one=1, two=2)""" + ) + + import builtins + import copy + import gc + import pickle + import types + import unittest + + class FrozenWithState(frozendict): + pass + + class FrozenWithSlot(frozendict): + __slots__ = ("extra",) + + class FrozenDictContractTests(unittest.TestCase): + def test_construction_identity_and_explicit_new(self): + original = frozendict(left=object()) + self.assertIs(frozendict(original), original) + explicit = frozendict.__new__(frozendict, original) + self.assertIsNot(explicit, original) + self.assertEqual(explicit, original) + self.assertIsNot(FrozenWithState(original), original) + self.assertIsNot(frozendict(FrozenWithState(original)), original) + original.__init__(object(), object(), ignored=object()) + self.assertEqual(list(original), ["left"]) + self.assertFalse(isinstance(original, dict)) + self.assertFalse(issubclass(frozendict, dict)) + + def test_call_ex_identity_preserves_explicit_type_call_distinction(self): + original = frozendict(answer=42) + self.assertIs(frozendict(*(original,)), original) + self.assertIs(frozendict(original, **{}), original) # noqa: PIE804 + self.assertIs(frozendict(*(original,), **{}), original) # noqa: PIE804 + self.assertIsNot(type.__call__(frozendict, original), original) + self.assertIsNot(frozendict.__call__(original), original) + + def test_unrelated_descriptors_cannot_mutate_or_access(self): + frozen = frozendict(protected=23) + for method, args in ( + (dict.__init__, ({"replacement": 8},)), + (dict.__setitem__, ("protected", 99)), + (dict.__delitem__, ("protected",)), + (dict.update, ({"replacement": 8},)), + (dict.clear, ()), + (dict.get, ("protected",)), + ): + with self.subTest(method=method), self.assertRaises(TypeError): + method(frozen, *args) + with self.assertRaises(TypeError): + frozendict.get({"protected": 23}, "protected") + self.assertEqual(frozen, {"protected": 23}) + for name in ( + "clear", + "update", + "pop", + "popitem", + "setdefault", + "__setitem__", + "__delitem__", + "__ior__", + ): + self.assertFalse(hasattr(frozen, name), name) + + def test_subclass_missing_is_only_a_subscript_hook(self): + class Lookup(frozendict): + def __missing__(self, key): + return "missing:" + key + + frozen = Lookup(present=23) + self.assertEqual(frozen["absent"], "missing:absent") + self.assertIsNone(frozen.get("absent")) + self.assertNotIn("absent", frozen) + + def test_fromkeys_constructor_calls_and_source_independence(self): + calls = [] + seed = frozendict(retained="old", replaced="old") + + class Factory(frozendict): + def __new__(cls, *args): + calls.append(args) + if not args: + return seed + return super().__new__(cls, *args) + + def __setitem__(self, key, value): + raise AssertionError("immutable construction called a mutator") + + frozen = Factory.fromkeys(iter(["replaced", "added", "replaced"]), "new") + self.assertIs(type(frozen), Factory) + self.assertEqual( + list(frozen.items()), + [("retained", "old"), ("replaced", "new"), ("added", "new")], + ) + self.assertEqual(seed, {"retained": "old", "replaced": "old"}) + self.assertEqual(len(calls), 2) + self.assertEqual(calls[0], ()) + self.assertIs(type(calls[1][0]), frozendict) + self.assertEqual(calls[1][0], frozen) + + def test_fromkeys_preserves_custom_constructor_result(self): + target = {"seed": 1} + + class MutableFactory(frozendict): + def __new__(cls): + return target + + result = MutableFactory.fromkeys(["new"], 2) + self.assertIs(result, target) + self.assertEqual(target, {"seed": 1, "new": 2}) + + def test_fromkeys_subclass_initialization_is_two_phase(self): + calls = [] + + class Initialized(frozendict): + def __init__(self, *args): + calls.append(args) + + result = Initialized.fromkeys(["first", "second"], 3) + self.assertIs(type(result), Initialized) + self.assertEqual(len(calls), 2) + self.assertEqual(calls[0], ()) + self.assertIs(type(calls[1][0]), frozendict) + self.assertEqual(calls[1][0], result) + + def test_fromkeys_constructs_before_rejecting_invalid_iterable(self): + calls = [] + + class Factory(frozendict): + def __new__(cls, *args): + calls.append(args) + return super().__new__(cls, *args) + + with self.assertRaises(TypeError): + Factory.fromkeys(123) + self.assertEqual(calls, [()]) + + def test_hash_uses_stored_key_hash_and_caches_values(self): + class Counted: + def __init__(self, value): + self.value = value + self.calls = 0 + + def __hash__(self): + self.calls += 1 + return self.value + + key, value = Counted(41), Counted(73) + frozen = frozendict([(key, value)]) + key_calls = key.calls + first = hash(frozen) + self.assertEqual(key.calls, key_calls) + self.assertEqual(value.calls, 1) + self.assertEqual(hash(frozen), first) + self.assertEqual(key.calls, key_calls) + self.assertEqual(value.calls, 1) + + def test_hash_failure_does_not_poison_cache(self): + class EventuallyHashable: + def __init__(self): + self.calls = 0 + + def __hash__(self): + self.calls += 1 + if self.calls == 1: + raise ValueError("not yet") + return 97 + + value = EventuallyHashable() + frozen = frozendict(retry=value) + with self.assertRaisesRegex(ValueError, "not yet"): + hash(frozen) + expected = hash(frozen) + self.assertEqual(hash(frozen), expected) + self.assertEqual(value.calls, 2) + + def test_mapping_equality_reuses_hashes_but_views_rehash(self): + class Key: + forbidden = False + + def __hash__(self): + if self.forbidden: + raise ValueError("rehash forbidden") + return 313 + + key = Key() + left = frozendict([(key, 7)]) + right = frozendict.__new__(frozendict, left) + mutable = dict(left) + key.forbidden = True + self.assertTrue(left == left) # noqa: PLR0124 + self.assertTrue(left.__eq__(left)) + self.assertTrue(left == right) + self.assertTrue(left == mutable) + self.assertTrue(mutable == left) + for method in ("keys", "items"): + with self.assertRaisesRegex(ValueError, "rehash forbidden"): + _ = getattr(left, method)() == getattr(right, method)() + + def test_value_hash_can_reenter_same_frozendict(self): + class Reentrant: + def __init__(self): + self.calls = 0 + self.owner = None + self.nested = None + + def __hash__(self): + self.calls += 1 + if self.calls == 1: + self.nested = hash(self.owner) + return 127 + + value = Reentrant() + frozen = frozendict(recursive=value) + value.owner = frozen + result = hash(frozen) + self.assertEqual(result, value.nested) + self.assertEqual(hash(frozen), result) + self.assertEqual(value.calls, 2) + + def test_union_result_type_order_and_identity(self): + frozen = frozendict(first=1, second=2) + result = frozen | {"first": 7, "last": 3} + self.assertIs(type(result), frozendict) + self.assertEqual( + list(result.items()), [("first", 7), ("second", 2), ("last", 3)] + ) + self.assertIs(type({"start": 4} | frozen), dict) + self.assertIs(frozen | {}, frozen) + self.assertIs(frozendict() | frozen, frozen) + child = FrozenWithState(frozen) + self.assertIs(type(child | {}), frozendict) + self.assertIsNot(child | {}, child) + before = frozen + frozen |= {"third": 3} + self.assertIsNot(frozen, before) + self.assertNotIn("third", before) + with self.assertRaises(TypeError): + frozen |= [("fourth", 4)] + + def test_explicit_reflected_union_uses_left_operand_type_and_order(self): + frozen = frozendict(left=1, shared="frozen") + mutable = {"right": 2, "shared": "mutable"} + result = dict.__ror__(mutable, frozen) + self.assertIs(type(result), frozendict) + self.assertEqual( + list(result.items()), [("left", 1), ("shared", "mutable"), ("right", 2)] + ) + result = frozendict.__ror__(frozen, mutable) + self.assertIs(type(result), dict) + self.assertEqual( + list(result.items()), [("right", 2), ("shared", "frozen"), ("left", 1)] + ) + + def test_frozen_globals_and_separate_writable_locals(self): + observed = [] + frozen = frozendict(answer=42, observed=observed, __builtins__=builtins) + local = {} + exec("observed.append(answer); local_answer = answer + 1", frozen, local) + self.assertEqual(observed, [42]) + self.assertEqual(local, {"local_answer": 43}) + self.assertEqual(eval("answer + 2", frozen), 44) + self.assertIs(eval("globals()", frozen), frozen) + self.assertNotIn("local_answer", frozen) + + def test_frozen_globals_require_existing_builtins(self): + for execute, source in ((eval, "40 + 2"), (exec, "pass")): + for frozen in (frozendict(answer=42), FrozenWithState(answer=42)): + with self.subTest(execute=execute, frozen_type=type(frozen)): + with self.assertRaisesRegex( + TypeError, + "cannot assign __builtins__ to frozendict globals", + ): + execute(source, frozen, {}) + self.assertNotIn("__builtins__", frozen) + + def test_functions_from_exec_keep_frozen_globals(self): + frozen = frozendict(answer=42, __builtins__=builtins) + local = {} + exec( + "def read():\n return answer\ndef assign():\n global answer\n answer = 99\ndef delete():\n global answer\n del answer", + frozen, + local, + ) + self.assertIs(local["read"].__globals__, frozen) + self.assertEqual(local["read"](), 42) + for name, suffix in (("assign", "assignment"), ("delete", "deletion")): + with self.assertRaises(TypeError) as caught: + local[name]() + self.assertEqual( + str(caught.exception), + "frozendict object does not support item " + suffix, + ) + with self.assertRaises(TypeError): + types.FunctionType(local["read"].__code__, frozen) + self.assertEqual(frozen["answer"], 42) + + def test_frozen_builtins_and_type_namespace(self): + frozen_builtins = frozendict(len=len) + frozen = frozendict(__builtins__=frozen_builtins) + self.assertEqual(eval("len((10, 20))", frozen), 2) + namespace = frozendict(value=17) + cls = type("FrozenNamespace", (), namespace) + self.assertEqual(cls.value, 17) + cls.value = 19 + self.assertEqual(namespace["value"], 17) + + def test_frozen_subclass_globals_keep_lookup_semantics(self): + class Namespace(frozendict): + def __getitem__(self, key): + if key == "answer": + return 81 + return super().__getitem__(key) + + frozen = Namespace(answer=42, __builtins__=builtins) + self.assertEqual(eval("answer", frozen), 81) + self.assertEqual(eval("(lambda: answer)()", frozen), 81) + self.assertEqual(eval("answer", frozen, {}), 42) + + def test_class_and_mapping_patterns(self): + frozen = frozendict(answer=42) + match frozen: + case frozendict(captured): + self.assertIs(captured, frozen) + case _: + self.fail("frozendict class pattern did not match") + match frozen: + case {"answer": captured, **remaining}: + self.assertEqual(captured, 42) + self.assertEqual(remaining, {}) + self.assertIs(type(remaining), dict) + case _: + self.fail("frozendict mapping pattern did not match") + + def test_hash_error_wrapper_preserves_exception_subclasses(self): + class SpecialTypeError(TypeError): + pass + + class BadKey: + error = TypeError + + def __hash__(self): + raise self.error("intentional hash failure") + + key = BadKey() + frozen = frozendict(present=1) + operations = ( + lambda: frozen[key], + lambda: frozen.get(key), + lambda: key in frozen, + lambda: frozendict([(key, 1)]), + lambda: frozendict.fromkeys([key]), + ) + for error in (TypeError, SpecialTypeError, ValueError): + key.error = error + for operation in operations: + with self.subTest(error=error, operation=operation): + with self.assertRaises(error) as caught: + operation() + self.assertIs(type(caught.exception), error) + expected = "intentional hash failure" + if error is TypeError: + name = BadKey.__qualname__ + if BadKey.__module__ not in ("builtins", "__main__"): + name = BadKey.__module__ + "." + name + expected = f"cannot use '{name}' as a frozendict key (intentional hash failure)" + self.assertEqual(str(caught.exception), expected) + + def test_copy_protocol_and_getnewargs_do_not_expose_storage(self): + frozen = frozendict(present=object()) + self.assertIs(copy.copy(frozen), frozen) + self.assertIs(frozen.copy(), frozen) + args = frozen.__getnewargs__() + self.assertIs(type(args), tuple) + self.assertIs(type(args[0]), dict) + args[0].clear() + self.assertEqual(list(frozen), ["present"]) + child = FrozenWithState(frozen) + child.extra = "state" + shallow = copy.copy(child) + self.assertIs(type(shallow), FrozenWithState) + self.assertEqual(shallow.extra, "state") + self.assertIs(type(child.copy()), frozendict) + deep = copy.deepcopy(frozen) + self.assertIsNot(deep, frozen) + + def test_subclass_copy_dispatch_tracks_iterator_override(self): + class NativeIteration(frozendict): + def keys(self): + return ["virtual"] + + def __getitem__(self, key): + return 88 + + class CustomIteration(NativeIteration): + def __iter__(self): + return iter(["virtual"]) + + native = NativeIteration(stored=7) + custom = CustomIteration(stored=7) + self.assertEqual(native.copy(), {"stored": 7}) + self.assertEqual(native | {"tail": 9}, {"stored": 7, "tail": 9}) + self.assertEqual(NativeIteration.fromkeys(["tail"], 9), {"tail": 9}) + self.assertEqual(custom.copy(), {"virtual": 88}) + self.assertEqual(custom | {"tail": 9}, {"virtual": 88, "tail": 9}) + self.assertEqual( + CustomIteration.fromkeys(["tail"], 9), {"virtual": 88, "tail": 9} + ) + + def test_deepcopy_recursive_aliases_from_both_roots(self): + loop = [] + frozen = frozendict(left=loop, right=loop) + loop.extend([frozen, frozen]) + for root in (frozen, loop): + result = copy.deepcopy(root) + copied_frozen = result if isinstance(result, frozendict) else result[0] + copied_loop = copied_frozen["left"] + self.assertIsNot(copied_frozen, frozen) + self.assertIsNot(copied_loop, loop) + self.assertIs(copied_frozen["right"], copied_loop) + self.assertIs(copied_loop[0], copied_frozen) + self.assertIs(copied_loop[1], copied_frozen) + + def test_pickle_recursive_payload_and_subclass_state(self): + for cls in (frozendict, FrozenWithState, FrozenWithSlot): + loop = [] + frozen = cls(left=loop, right=loop) + loop.extend([frozen, frozen]) + if cls is not frozendict: + frozen.extra = loop + for protocol in range(pickle.HIGHEST_PROTOCOL + 1): + with self.subTest(cls=cls, protocol=protocol): + if protocol < 2: + with self.assertRaises(TypeError): + pickle.dumps(frozen, protocol) + continue + result = pickle.loads(pickle.dumps(frozen, protocol)) + self.assertIs(type(result), cls) + self.assertIsNot(result, frozen) + self.assertIs(result["left"], result["right"]) + self.assertIs(result["left"][0], result) + self.assertIs(result["left"][1], result) + if cls is not frozendict: + self.assertIs(result.extra, result["left"]) + + def test_gc_does_not_expose_mutable_backing_mapping(self): + marker = object() + frozen = frozendict(private_payload=marker) + for referent in gc.get_referents(frozen): + self.assertFalse(isinstance(referent, dict)) + self.assertIs(frozen["private_payload"], marker) + + def test_gc_does_not_publish_partly_constructed_mapping(self): + marker = object() + observed = [] + + def source(): + yield "private_payload", marker + observed.extend( + candidate + for candidate in gc.get_objects() + if type(candidate) is frozendict + and candidate.get("private_payload") is marker + ) + yield "finished", True + + frozen = frozendict(source()) + self.assertEqual(observed, []) + self.assertIs(frozen["private_payload"], marker) + self.assertIs(frozen["finished"], True) + + import gc + import operator + import unittest + import weakref + + class FrozenDictEdgeCaseTests(unittest.TestCase): + def test_hash_matches_frozenset_of_items(self): + cases = ( + frozendict(), + frozendict([(1, 2), (-3, 4), (2**130, -7)]), + frozendict(first=1, second=2), + frozendict(second=2, first=1), + frozendict([(1, "one"), ("two", 2), (("three",), frozenset({3}))]), + ) + for frozen in cases: + with self.subTest(frozen=frozen): + self.assertEqual(hash(frozen), hash(frozenset(frozen.items()))) + + def test_gc_cycles_and_class_reassignment_cannot_expose_mutation(self): + frozen = frozendict(protected=23) + original_hash = hash(frozen) + with self.assertRaises(TypeError): + frozen.__class__ = dict + self.assertIs(type(frozen), frozendict) + self.assertEqual(frozen, {"protected": 23}) + self.assertEqual(hash(frozen), original_hash) + + class Marker: + pass + + def make_cycle(include_views): + marker = Marker() + loop = [marker] + owner = frozendict(loop=loop) + loop.append(owner) + if include_views: + loop.extend( + (owner.keys(), owner.values(), owner.items(), iter(owner)) + ) + return weakref.ref(marker) + + for include_views in (False, True): + with self.subTest(include_views=include_views): + marker_ref = make_cycle(include_views) + gc.collect() + self.assertIsNone(marker_ref()) + + def test_union_subclass_dispatch_matrix(self): + for base in (dict, frozendict): + + class NativeIteration(base): + def keys(self): + return ["virtual"] + + def __getitem__(self, key): + return 88 + + class CustomIteration(NativeIteration): + def __iter__(self): + return iter(["virtual"]) + + for cls in (NativeIteration, CustomIteration): + custom = cls is CustomIteration + for storage in ({}, {"stored": 7}): + subject = cls(storage) + copied = {"virtual": 88} if custom and storage else storage + merged = {"virtual": 88} if custom else storage + cases = ( + ("left_dict", lambda: subject | {}, base, copied), + ( + "left_frozen", + lambda: subject | frozendict(), + base, + copied, + ), + ("right_dict", lambda: {} | subject, dict, merged), + ( + "right_frozen", + lambda: frozendict() | subject, + frozendict, + merged, + ), + ( + "construction", + lambda: frozendict(subject), + frozendict, + merged, + ), + ) + for name, operation, expected_type, expected in cases: + with self.subTest( + base=base, cls=cls, storage=storage, operation=name + ): + result = operation() + self.assertIs(type(result), expected_type) + self.assertEqual(result, expected) + self.assertEqual( + list(subject.items()), list(storage.items()) + ) + + if base is frozendict: + with self.subTest( + base=base, cls=cls, storage=storage, operation="copy" + ): + result = subject.copy() + self.assertIs(type(result), frozendict) + self.assertIsNot(result, subject) + self.assertEqual(result, copied) + + def test_mapping_equality_uses_left_value_and_stored_key_hash(self): + trace = [] + + class Key: + def __init__(self, side): + self.side = side + + def __hash__(self): + trace.append(("hash", self.side)) + return 313 + + def __eq__(self, other): + trace.append(("key", self.side, other.side)) + return True + + class Value: + def __init__(self, side): + self.side = side + + def __eq__(self, other): + trace.append(("value", self.side, other.side)) + return self.side == "left" + + for left_type in (dict, frozendict): + for right_type in (dict, frozendict): + left = left_type([(Key("left"), Value("left"))]) + right = right_type([(Key("right"), Value("right"))]) + with self.subTest(left_type=left_type, right_type=right_type): + trace.clear() + self.assertTrue(left == right) + self.assertEqual( + trace, + [("key", "right", "left"), ("value", "left", "right")], + ) + trace.clear() + self.assertFalse(right == left) + self.assertEqual( + trace, + [("key", "left", "right"), ("value", "right", "left")], + ) + + def test_view_comparison_direction_and_asymmetric_values(self): + trace = [] + + class Key: + def __init__(self, side): + self.side = side + + def __hash__(self): + trace.append(("hash", self.side)) + return 313 + + def __eq__(self, other): + trace.append(("key", self.side, other.side)) + return True + + class Value: + def __init__(self, side): + self.side = side + + def __eq__(self, other): + trace.append(("value", self.side, other.side)) + return self.side == "left" + + for left_type in (dict, frozendict): + for right_type in (dict, frozendict): + left = left_type([(Key("left"), Value("left"))]) + right = right_type([(Key("right"), Value("right"))]) + for method in ("keys", "items"): + for op in ( + operator.eq, + operator.ne, + operator.lt, + operator.le, + operator.gt, + operator.ge, + ): + with self.subTest( + left_type=left_type, + right_type=right_type, + method=method, + operation=op.__name__, + ): + trace.clear() + result = op( + getattr(left, method)(), getattr(right, method)() + ) + if op in (operator.lt, operator.gt): + self.assertFalse(result) + self.assertEqual(trace, []) + continue + reverse = op is operator.ge + needle, container = ( + ("right", "left") if reverse else ("left", "right") + ) + expected_trace = [ + ("hash", needle), + ("key", container, needle), + ] + contained = True + if method == "items": + expected_trace.append(("value", container, needle)) + contained = container == "left" + self.assertEqual( + result, + not contained if op is operator.ne else contained, + ) + self.assertEqual(trace, expected_trace) + + def test_view_hash_error_context_comes_from_containment_mapping(self): + class Key: + broken = False + + def __hash__(self): + if self.broken: + raise TypeError("intentional view hash error") + return 313 + + for left_type in (dict, frozendict): + for right_type in (dict, frozendict): + key = Key() + left = left_type([(key, 1)]) + right = right_type([(key, 1)]) + key.broken = True + for method in ("keys", "items"): + for op in (operator.eq, operator.ne, operator.le, operator.ge): + with self.subTest( + left_type=left_type, + right_type=right_type, + method=method, + operation=op.__name__, + ): + container = ( + left_type if op is operator.ge else right_type + ) + with self.assertRaises(TypeError) as caught: + op( + getattr(left, method)(), + getattr(right, method)(), + ) + self.assertTrue( + str(caught.exception).endswith( + f"as a {container.__name__} key (intentional view hash error)" + ), + str(caught.exception), + ) + + def test_view_equality_checks_mutation_of_the_iterated_mapping(self): + for op in (operator.eq, operator.le, operator.ge): + left = {"key": object()} + + class Value: + def __eq__(self, other): + left.clear() + return True + + right = frozendict(key=Value()) + with self.subTest(operation=op.__name__): + if op is operator.ge: + self.assertTrue(op(left.items(), right.items())) + else: + with self.assertRaisesRegex( + RuntimeError, "dictionary changed size during iteration" + ): + op(left.items(), right.items()) + self.assertEqual(left, {}) + + def test_fromkeys_merges_constructor_result_before_iterable_validation(self): + calls = [] + + class Source(frozendict): + def __iter__(self): + return iter(()) + + def keys(self): + calls.append("keys") + return [] + + source = Source() + + class Factory(frozendict): + def __new__(cls, *args): + calls.append("new") + return source + + with self.assertRaises(TypeError): + Factory.fromkeys(123) + self.assertEqual(calls, ["new", "keys"]) + self.assertEqual(source, {}) + + def test_fromkeys_source_error_precedes_invalid_iterable(self): + class Source(frozendict): + def __iter__(self): + return iter(()) + + def keys(self): + raise ValueError("source merge failed") + + source = Source() + + class Factory(frozendict): + def __new__(cls, *args): + return source + + with self.assertRaisesRegex(ValueError, "source merge failed"): + Factory.fromkeys(123) + + import builtins + import gc + import sys + import types + import unittest + + class FrozenGlobalsTests(unittest.TestCase): + def setUp(self): + self.ns = frozendict( + __builtins__=builtins.__dict__, + __name__="frozen_test", + answer=42, + sys=sys, + ) + + def test_eval_and_implicit_scope(self): + self.assertIs(eval("globals()", self.ns), self.ns) + self.assertIs(eval("locals()", self.ns), self.ns) + self.assertEqual(eval("eval('answer')", self.ns), 42) + self.assertEqual(eval("(lambda: answer)()", self.ns), 42) + + def test_function_and_materialized_frame(self): + local = {} + exec( + "def f():\n return sys._getframe(), globals(), answer\n", self.ns, local + ) + f = local["f"] + self.assertIs(f.__globals__, self.ns) + self.assertEqual(f.__module__, "frozen_test") + for _ in range(300): + frame, ns, answer = f() + self.assertIs(frame.f_globals, self.ns) + self.assertIs(ns, self.ns) + self.assertEqual(answer, 42) + gc.collect() + self.assertIs(frame.f_globals, self.ns) + + def test_generator_frames(self): + local = {} + exec( + "def gen():\n yield sys._getframe(), globals(), answer\n yield answer\n", + self.ns, + local, + ) + gen = local["gen"]() + self.assertIs(gen.gi_frame.f_globals, self.ns) + frame, ns, answer = next(gen) + self.assertIs(frame.f_globals, self.ns) + self.assertIs(ns, self.ns) + self.assertEqual(answer, 42) + self.assertEqual(next(gen), 42) + with self.assertRaises(StopIteration): + next(gen) + gc.collect() + self.assertIs(frame.f_globals, self.ns) + + def test_coroutine_frames(self): + local = {} + exec( + "async def coro():\n return sys._getframe(), globals(), answer\n", + self.ns, + local, + ) + coro = local["coro"]() + self.assertIs(coro.cr_frame.f_globals, self.ns) + with self.assertRaises(StopIteration) as stopped: + coro.send(None) + frame, ns, answer = stopped.exception.value + self.assertIs(frame.f_globals, self.ns) + self.assertIs(ns, self.ns) + self.assertEqual(answer, 42) + + def test_async_generator_frames(self): + local = {} + exec( + "async def agen():\n yield sys._getframe(), globals(), answer\n", + self.ns, + local, + ) + agen = local["agen"]() + self.assertIs(agen.ag_frame.f_globals, self.ns) + with self.assertRaises(StopIteration) as stopped: + agen.asend(None).send(None) + frame, ns, answer = stopped.exception.value + self.assertIs(frame.f_globals, self.ns) + self.assertIs(ns, self.ns) + self.assertEqual(answer, 42) + with self.assertRaises(StopIteration): + agen.aclose().send(None) + + def test_local_namespace_and_imports(self): + local = {} + exec( + "import math\nclass C:\n value = answer\ndef f():\n import math\n return math.sqrt(answer)\n", + self.ns, + local, + ) + self.assertEqual(local["C"].value, 42) + self.assertEqual(local["C"].__module__, "frozen_test") + self.assertEqual(local["f"](), local["math"].sqrt(42)) + self.assertNotIn("math", self.ns) + + def test_missing_builtins(self): + for func, source in [(eval, ""), (exec, "")]: + with self.assertRaises(TypeError) as caught: + func(source, frozendict()) + self.assertEqual( + str(caught.exception), + "cannot assign __builtins__ to frozendict globals", + ) + + def test_writes_and_deletes(self): + for source, error in [ + ("answer = 0", "'frozendict' object does not support item assignment"), + ( + "global answer; answer = 0", + "frozendict object does not support item assignment", + ), + ( + "global answer; del answer", + "frozendict object does not support item deletion", + ), + ( + "global missing; del missing", + "frozendict object does not support item deletion", + ), + ]: + with self.assertRaises(TypeError) as caught: + exec(source, self.ns) + self.assertEqual(str(caught.exception), error) + self.assertEqual(self.ns["answer"], 42) + with self.assertRaises(NameError): + exec("del answer", self.ns) + + def test_explicit_function_constructor_stays_dict_only(self): + with self.assertRaises(TypeError): + types.FunctionType((lambda: 42).__code__, self.ns) + + def test_subclass_name_and_global_lookup(self): + class F(frozendict): + def __getitem__(self, key): + if key == "answer": + return 81 + return super().__getitem__(key) + + ns = F(self.ns) + self.assertEqual(eval("answer", ns), 81) + self.assertEqual(eval("(lambda: answer)()", ns), 81) + self.assertEqual(eval("answer", ns, {}), 42) + + def test_subclass_function_metadata_reads_storage(self): + class F(frozendict): + def __getitem__(self, key): + if key == "__name__": + return "wrong_name" + if key == "__builtins__": + return {"len": lambda x: 99} + return super().__getitem__(key) + + ns = F(self.ns) + f = eval("lambda: len([])", ns) + self.assertEqual(f.__module__, "frozen_test") + self.assertIs(f.__globals__, ns) + self.assertEqual(f(), 0) + + def test_subclass_builtins_assignment_hook(self): + calls = [] + + class F(frozendict): + def __setitem__(self, key, value): + calls.append((key, value)) + + for operation, source, mode, expected in ( + (eval, "1", "eval", 1), + (exec, "pass", "exec", None), + ): + namespace = F() + calls.clear() + self.assertEqual( + operation(compile(source, "", mode), namespace), + expected, + ) + self.assertEqual(calls, [("__builtins__", builtins.__dict__)]) + self.assertNotIn("__builtins__", namespace) + calls.clear() + with self.assertRaisesRegex( + TypeError, "frozendict object does not support item assignment" + ): + operation(source, namespace) + self.assertEqual(calls, [("__builtins__", builtins.__dict__)]) + + def test_implicit_frozen_builtins_namespace_counts(self): + import _testinternalcapi + + function = eval( + "lambda: len(())", + frozendict(__builtins__=frozendict(len=len)), + ) + counts = _testinternalcapi.get_code_var_counts(function) + self.assertEqual( + counts["unbound"]["globals"], + {"total": 1, "numglobal": 0, "numbuiltin": 1, "numunknown": 0}, + ) + for argument in ("globalsns", "builtinsns"): + with self.assertRaises(TypeError): + _testinternalcapi.get_code_var_counts( + function, **{argument: frozendict()} + ) + + def test_warning_context_reads_frozen_storage(self): + import warnings + + class F(frozendict): + def __getitem__(self, key): + if key in ("__warningregistry__", "__name__"): + raise AssertionError("warning context must read stored entries") + return super().__getitem__(key) + + for has_registry in (False, True): + entries = dict(self.ns, warn=warnings.warn) + if has_registry: + entries["__warningregistry__"] = {} + function = eval("lambda: warn('probe')", F(entries)) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + if has_registry: + function() + self.assertEqual(len(caught), 1) + self.assertEqual(str(caught[0].message), "probe") + else: + with self.assertRaisesRegex( + TypeError, + "frozendict object does not support item assignment", + ): + function() + self.assertEqual(caught, []) + + def test_frozen_builtins(self): + ns = frozendict(__builtins__=frozendict({"len": len})) + self.assertEqual(eval("(lambda: len([1, 2]))()", ns), 2) + + def test_shared_code_does_not_keep_mutable_global_specialization(self): + code = compile("lambda: answer", "", "eval") + mutable = {"__builtins__": builtins.__dict__, "answer": 19} + f = eval(code, mutable) + for _ in range(300): + self.assertEqual(f(), 19) + frozen_f = eval(code, self.ns) + for _ in range(300): + self.assertEqual(frozen_f(), 42) + self.assertEqual(f(), 19) + mutable["answer"] = 23 + self.assertEqual(f(), 23) + self.assertEqual(frozen_f(), 42) + + import types + import unittest + + class FrozenTypeNamespaceTests(unittest.TestCase): + def test_exact_frozen_namespace_is_copied(self): + namespace = frozendict(value=17, __qualname__="Outer.FrozenNamespace") + original_hash = hash(namespace) + cls = type("FrozenNamespace", (), namespace) + self.assertEqual(cls.value, 17) + self.assertEqual(cls.__qualname__, "Outer.FrozenNamespace") + cls.value = 19 + cls.added = 23 + del cls.value + self.assertEqual( + namespace, + {"value": 17, "__qualname__": "Outer.FrozenNamespace"}, + ) + self.assertEqual(hash(namespace), original_hash) + self.assertNotIn("__module__", namespace) + self.assertNotIn("added", namespace) + + def test_builtin_iterator_bypasses_subclass_mapping_overrides(self): + for base in (dict, frozendict): + + class Namespace(base): + def keys(self): + raise AssertionError("keys must not be called") + + def __getitem__(self, key): + raise AssertionError("__getitem__ must not be called") + + def __len__(self): + raise AssertionError("__len__ must not be called") + + with self.subTest(base=base): + namespace = Namespace(value=17, __qualname__="Outer.C") + cls = type("C", (), namespace) + self.assertEqual(cls.value, 17) + self.assertEqual(cls.__qualname__, "Outer.C") + self.assertEqual(base.__getitem__(namespace, "value"), 17) + + def test_custom_iterator_uses_mapping_order_and_values(self): + for base in (dict, frozendict): + calls = [] + + class Namespace(base): + def __iter__(self): + raise AssertionError("the copy must use keys, not __iter__") + + def keys(self): + calls.append("keys") + return ["second", "first"] + + def __getitem__(self, key): + calls.append(key) + return {"first": 11, "second": 22}[key] + + with self.subTest(base=base): + namespace = Namespace(stored=3) + cls = type("C", (), namespace) + self.assertEqual(calls, ["keys", "second", "first"]) + self.assertEqual((cls.first, cls.second), (11, 22)) + self.assertFalse(hasattr(cls, "stored")) + self.assertEqual( + [key for key in cls.__dict__ if not key.startswith("__")], + ["second", "first"], + ) + self.assertEqual(base.__getitem__(namespace, "stored"), 3) + + def test_empty_subclass_copy_bypasses_mapping_overrides(self): + for base in (dict, frozendict): + + class Namespace(base): + def __iter__(self): + raise AssertionError("__iter__ must not be called") + + def keys(self): + raise AssertionError("keys must not be called") + + def __getitem__(self, key): + raise AssertionError("__getitem__ must not be called") + + def __len__(self): + return 1 + + with self.subTest(base=base): + namespace = Namespace() + cls = type("C", (), namespace) + self.assertEqual(cls.__name__, "C") + self.assertEqual(base.__len__(namespace), 0) + + def test_class_cells_use_a_mutable_copy(self): + classcell = types.CellType() + dictcell = types.CellType() + namespace = frozendict( + value=17, + __classcell__=classcell, + __classdictcell__=dictcell, + ) + cls = type("C", (), namespace) + self.assertIs(classcell.cell_contents, cls) + classdict = dictcell.cell_contents + self.assertIs(type(classdict), dict) + self.assertIsNot(classdict, namespace) + self.assertNotIn("__classcell__", classdict) + self.assertNotIn("__classdictcell__", classdict) + self.assertIs(namespace["__classcell__"], classcell) + self.assertIs(namespace["__classdictcell__"], dictcell) + cls.value = 19 + self.assertEqual(classdict["value"], 19) + self.assertEqual(namespace["value"], 17) + + def test_descriptor_initialization_does_not_change_source(self): + calls = [] + + class Descriptor: + def __set_name__(self, owner, name): + calls.append((owner, name)) + owner.added = 23 + + descriptor = Descriptor() + namespace = frozendict(field=descriptor) + cls = type("C", (), namespace) + self.assertEqual(calls, [(cls, "field")]) + self.assertEqual(cls.added, 23) + self.assertIs(namespace["field"], descriptor) + self.assertEqual(list(namespace), ["field"]) + + def test_winning_metaclass_receives_original_namespace(self): + for base in (dict, frozendict): + + class Namespace(base): + def __iter__(self): + raise AssertionError("namespace must not be copied") + + def keys(self): + raise AssertionError("namespace must not be copied") + + namespace = Namespace(value=17) + + class Meta(type): + def __new__(metacls, name, bases, ns): + if name == "Child": + return ns + return super().__new__(metacls, name, bases, ns) + + class Base(metaclass=Meta): + pass + + with self.subTest(base=base): + self.assertIs(type("Child", (Base,), namespace), namespace) + + def test_keywords_and_explicit_type_new(self): + calls = [] + + class Base: + def __init_subclass__(cls, *, marker): + calls.append((cls, marker)) + + namespace = frozendict(value=17) + cls = type("C", (Base,), namespace, marker=23) + explicit = type.__new__(type, "D", (Base,), namespace, marker=29) + self.assertEqual(calls, [(cls, 23), (explicit, 29)]) + self.assertEqual((cls.value, explicit.value), (17, 17)) + self.assertEqual(namespace, {"value": 17}) + + def test_non_dictionary_mappings_remain_rejected(self): + class Mapping: + def keys(self): + return ["value"] + + def __getitem__(self, key): + return 17 + + for namespace in (types.MappingProxyType({}), Mapping(), [], ()): + with self.subTest(namespace=namespace), self.assertRaises(TypeError): + type("C", (), namespace) + + frozen_suite = unittest.TestSuite( + unittest.defaultTestLoader.loadTestsFromTestCase(case) + for case in ( + FrozenDictContractTests, + FrozenDictEdgeCaseTests, + FrozenGlobalsTests, + FrozenTypeNamespaceTests, + ) + ) + assert frozen_suite.countTestCases() == 61 + frozen_result = unittest.TextTestRunner(verbosity=2).run(frozen_suite) + assert frozen_result.wasSuccessful() + assert not frozen_result.skipped + + +# Final 3.15 rewrites exact hash TypeError without an explicit cause, and +# propagates formatting failures instead of replacing them with diagnostics. +def check_mapping_hash_error_context(): + class SpecialTypeError(TypeError): + pass + + for mapping_type in (dict, frozendict): + for case in ( + "plain", + "empty", + "multiple", + "bad_str", + "bad_repr", + "subclass", + "subclass_bad_str", + "list", + ): + for handled in (False, True): + for operation in ("getitem", "get", "contains"): + callback_error = RuntimeError("format failed") + callbacks = [] + + class BadText: + def __str__(self): + callbacks.append("str") + raise callback_error + + def __repr__(self): + callbacks.append("repr") + raise callback_error + + args = { + "plain": ("hash failed",), + "empty": (), + "multiple": ("one", 2), + "bad_str": (BadText(),), + "bad_repr": ("one", BadText()), + "subclass": ("hash failed",), + "subclass_bad_str": (BadText(),), + "list": (), + }[case] + error_type = ( + SpecialTypeError if case.startswith("subclass") else TypeError + ) + original = error_type(*args) + outer = ValueError("outer") + + class Key: + def __hash__(self): + raise original + + key = [] if case == "list" else Key() + mapping = mapping_type() + + def invoke(): + try: + if operation == "getitem": + mapping[key] + elif operation == "get": + mapping.get(key) + else: + key in mapping + except Exception as error: + assert error.__cause__ is None + assert error.__context__ is (outer if handled else None) + assert error.__suppress_context__ is False + if case.startswith("subclass"): + assert error is original + assert callbacks == [] + elif case in ("bad_str", "bad_repr"): + assert error is callback_error + assert callbacks == [ + "str" if case == "bad_str" else "repr" + ] + else: + assert type(error) is TypeError + assert error is not original + name = "list" if case == "list" else Key.__qualname__ + message = { + "plain": "hash failed", + "empty": "", + "multiple": "('one', 2)", + "list": "unhashable type: 'list'", + }[case] + assert str(error) == ( + f"cannot use '{name}' as a {mapping_type.__name__} key ({message})" + ) + else: + raise AssertionError("hash failure was not propagated") + + if handled: + try: + raise outer + except ValueError: + invoke() + else: + invoke() + + +check_mapping_hash_error_context() diff --git a/scripts/update_lib/__main__.py b/scripts/update_lib/__main__.py index 49399db6f43..7eb91970f2f 100644 --- a/scripts/update_lib/__main__.py +++ b/scripts/update_lib/__main__.py @@ -8,6 +8,7 @@ python scripts/update_lib migrate cpython/Lib/test/test_foo.py python scripts/update_lib patches --from Lib/test/foo.py --to cpython/Lib/test/foo.py python scripts/update_lib auto-mark Lib/test/test_foo.py + python scripts/update_lib copy-test SOURCE DESTINATION --test ClassName.method_name --source-sha256 HASH --destination-sha256 HASH """ import argparse @@ -49,6 +50,11 @@ def main(argv: list[str] | None = None) -> int: help="Copy library file/directory from CPython (delete existing first)", add_help=False, ) + subparsers.add_parser( + "copy-test", + help="Copy one complete canonical test from hash-pinned inputs", + add_help=False, + ) subparsers.add_parser( "deps", help="Show dependency information for a module", @@ -72,6 +78,11 @@ def main(argv: list[str] | None = None) -> int: return copy_lib_main(remaining) + if args.command == "copy-test": + from update_lib.cmd_copy_test import main as copy_test_main + + return copy_test_main(remaining) + if args.command == "migrate": from update_lib.cmd_migrate import main as migrate_main diff --git a/scripts/update_lib/cmd_copy_test.py b/scripts/update_lib/cmd_copy_test.py new file mode 100644 index 00000000000..5e14959bb02 --- /dev/null +++ b/scripts/update_lib/cmd_copy_test.py @@ -0,0 +1,134 @@ +"""Copy one complete canonical test definition without refreshing its module.""" + +import argparse +import ast +import hashlib +import json +from pathlib import Path + +from update_lib.patch_spec import _iter_class_methods, _iter_module_classes + + +def _definition(contents: str, qualified_name: str): + parts = qualified_name.split(".") + if len(parts) != 2 or not all(part.isidentifier() for part in parts): + raise ValueError("test name must be ClassName.method_name") + tree = ast.parse(contents) + classes = [node for node in _iter_module_classes(tree) if node.name == parts[0]] + if len(classes) != 1: + raise ValueError(f"Expected one class named {parts[0]}, found {len(classes)}") + methods = [ + node for node in _iter_class_methods(classes[0]) if node.name == parts[1] + ] + if len(methods) != 1: + raise ValueError( + f"Expected one method named {qualified_name}, found {len(methods)}" + ) + method = methods[0] + parents = { + child: parent + for parent in ast.walk(tree) + for child in ast.iter_child_nodes(parent) + } + guards = [] + child = method + while child in parents: + parent = parents[child] + if isinstance(parent, ast.If): + branch = "body" if child in parent.body else "orelse" + guards.append((ast.dump(parent.test), branch)) + elif not isinstance(parent, (ast.Module, ast.ClassDef)): + raise ValueError("Only module/class scopes and if guards are supported") + child = parent + start = min([method.lineno, *(node.lineno for node in method.decorator_list)]) - 1 + return method, guards, start, method.end_lineno + + +def copy_test(source: str, destination: str, qualified_name: str) -> str: + """Replace an unambiguous whole definition, preserving all surrounding bytes. + + Guard, signature, decorator and indentation differences require a broader + migration rather than silently changing the selected test's execution context. + """ + src, src_guards, src_start, src_end = _definition(source, qualified_name) + dst, dst_guards, dst_start, dst_end = _definition(destination, qualified_name) + if src_guards != dst_guards: + raise ValueError("Source and destination test guards differ") + + def signature(node): + return ( + ast.dump(node.args), + ast.dump(node.returns) if node.returns is not None else None, + [ast.dump(param) for param in getattr(node, "type_params", [])], + ) + + if type(src) is not type(dst) or signature(src) != signature(dst): + raise ValueError("Source and destination test signatures differ") + if [ast.dump(node) for node in src.decorator_list] != [ + ast.dump(node) for node in dst.decorator_list + ]: + raise ValueError( + "Source and destination decorators differ; preserve markers explicitly" + ) + if src.col_offset != dst.col_offset: + raise ValueError("Source and destination indentation differs") + source_lines = source.splitlines(keepends=True) + destination_lines = destination.splitlines(keepends=True) + # AST equality omits comments, including project marker rationales. Require + # the complete decorator block to match before replacing the definition. + if ( + source_lines[src_start : src.lineno - 1] + != destination_lines[dst_start : dst.lineno - 1] + ): + raise ValueError( + "Source and destination decorator text differs; preserve markers explicitly" + ) + definition = source_lines[src_start:src_end] + result = "".join( + destination_lines[:dst_start] + definition + destination_lines[dst_end:] + ) + _, _, result_start, result_end = _definition(result, qualified_name) + if result.splitlines(keepends=True)[result_start:result_end] != definition: + raise ValueError("Copied definition is not byte-identical to the source") + return result + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("source", type=Path) + parser.add_argument("destination", type=Path) + parser.add_argument("--test", required=True, help="Exact ClassName.method_name") + parser.add_argument("--source-sha256", required=True) + parser.add_argument("--destination-sha256", required=True) + args = parser.parse_args(argv) + try: + source = args.source.read_bytes() + destination = args.destination.read_bytes() + for label, contents, expected in ( + ("source", source, args.source_sha256), + ("destination", destination, args.destination_sha256), + ): + if hashlib.sha256(contents).hexdigest() != expected: + raise ValueError(f"{label} SHA256 does not match the pinned input") + result = copy_test( + source.decode("utf-8"), destination.decode("utf-8"), args.test + ).encode("utf-8") + # Recheck the destination immediately before writing so stale invocations + # fail instead of overwriting an intervening edit. + if args.destination.read_bytes() != destination: + raise ValueError("Destination changed during test preparation") + args.destination.write_bytes(result) + except (OSError, UnicodeError, SyntaxError, ValueError) as exc: + parser.exit(1, f"Error: {exc}\n") + print( + json.dumps( + { + "test": args.test, + "source_sha256": args.source_sha256, + "destination_before_sha256": args.destination_sha256, + "destination_after_sha256": hashlib.sha256(result).hexdigest(), + }, + sort_keys=True, + ) + ) + return 0 diff --git a/scripts/update_lib/patch_spec.py b/scripts/update_lib/patch_spec.py index 996b78bfa18..994fc53b261 100644 --- a/scripts/update_lib/patch_spec.py +++ b/scripts/update_lib/patch_spec.py @@ -197,7 +197,10 @@ def visit_ClassDef(self, node: ast.ClassDef): patch = patch._replace(test_name="__self__") self.patches.append(patch) - self.generic_visit(node) + # Keep conditional methods in this namespace, but do not enter + # nested classes or the bodies of methods. + for item in _iter_class_methods(node): + self.visit(item) class PatchEntry(typing.NamedTuple): @@ -228,22 +231,35 @@ def iter_patch_entries( yield from visitor.patches -def iter_tests( - tree: ast.Module, -) -> "Iterator[tuple[ast.ClassDef, ast.FunctionDef | ast.AsyncFunctionDef]]": - for key, nodes in ast.iter_fields(tree): - if key != "body": - continue +def _iter_module_classes(node: ast.AST) -> "Iterator[ast.ClassDef]": + """Find module-scope classes, including conditional definitions. - for cls_node in nodes: - if not isinstance(cls_node, ast.ClassDef): - continue + Function-local and class-nested definitions have different namespaces and + cannot be addressed by the patch format's unqualified class names. + """ + for child in ast.iter_child_nodes(node): + if isinstance(child, ast.ClassDef): + yield child + elif not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + yield from _iter_module_classes(child) - for fn_node in cls_node.body: - if not isinstance(fn_node, (ast.FunctionDef, ast.AsyncFunctionDef)): - continue - yield (cls_node, fn_node) +def _iter_class_methods( + node: ast.AST, +) -> "Iterator[ast.FunctionDef | ast.AsyncFunctionDef]": + for child in ast.iter_child_nodes(node): + if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)): + yield child + elif not isinstance(child, (ast.ClassDef, ast.Lambda)): + yield from _iter_class_methods(child) + + +def iter_tests( + tree: ast.Module, +) -> "Iterator[tuple[ast.ClassDef, ast.FunctionDef | ast.AsyncFunctionDef]]": + for cls_node in _iter_module_classes(tree): + for fn_node in _iter_class_methods(cls_node): + yield (cls_node, fn_node) def iter_patches(contents: str) -> "Iterator[PatchEntry]": @@ -286,24 +302,28 @@ def _iter_patch_lines( # Track class bases for inherited async method lookup class_bases: dict[str, list[str]] = {} all_classes = set() - all_class_nodes = [] - for node in tree.body: - if not isinstance(node, ast.ClassDef): - continue - + all_class_nodes = list(_iter_module_classes(tree)) + for node in all_class_nodes: + if node.name in all_classes: + raise ValueError(f"Ambiguous module-scope class name: {node.name}") all_classes.add(node.name) - all_class_nodes.append(node) + method_counts = collections.Counter( + item.name for item in _iter_class_methods(node) + ) + for method in patches.get(node.name, {}): + if method_counts[method] > 1: + raise ValueError(f"Ambiguous method name: {node.name}.{method}") - for node in tree.body: + for node in all_class_nodes: if isinstance(node, ast.ClassDef): - cache[node.name] = node.end_lineno + cache[node.name] = (node.end_lineno, node.col_offset + len(DEFAULT_INDENT)) class_bases[node.name] = [ base.id for base in node.bases if isinstance(base, ast.Name) and base.id in all_classes ] cls_async: set[str] = set() - for item in node.body: + for item in _iter_class_methods(node): if isinstance(item, ast.AsyncFunctionDef): cls_async.add(item.name) if cls_async: @@ -319,10 +339,11 @@ def _iter_patch_lines( # Phase 2: Iterate and mark inherited tests for cls_name, tests in sorted(patches.items()): - lineno = cache.get(cls_name) - if not lineno: + location = cache.get(cls_name) + if not location: print(f"WARNING: {cls_name} does not exist in remote file", file=sys.stderr) continue + lineno, indent = location for test_name, specs in sorted(tests.items()): if test_name == "__self__": @@ -356,7 +377,7 @@ async def {test_name}(self): def {test_name}(self): {DEFAULT_INDENT}return super().{test_name}() """.rstrip() - yield (lineno, textwrap.indent(patch_lines, DEFAULT_INDENT)) + yield (lineno, textwrap.indent(patch_lines, " " * indent)) # Phase 3: Mark the class itself for cls_node in all_class_nodes: diff --git a/scripts/update_lib/tests/test_copy_test.py b/scripts/update_lib/tests/test_copy_test.py new file mode 100644 index 00000000000..5dc61452c90 --- /dev/null +++ b/scripts/update_lib/tests/test_copy_test.py @@ -0,0 +1,206 @@ +"""Tests for bounded, hash-pinned canonical test copying.""" + +import ast +import contextlib +import hashlib +import io +import tempfile +import unittest +from pathlib import Path + +from update_lib.cmd_copy_test import copy_test, main + + +class TestCopyTest(unittest.TestCase): + destination = '''# Destination-only header +class Finder: + if int.__doc__: + def example(): + r"""Count is below 800. π\\n""" + + def untouched(self): + return "destination" +''' + source = '''# Canonical header must not be copied +class Finder: + if int.__doc__: + def example(): + r"""Count is below 850. π\\n""" + + def untouched(self): + return "canonical" +''' + + def test_complete_definition_and_surroundings(self): + result = copy_test(self.source, self.destination, "Finder.example") + self.assertEqual(result, self.destination.replace("below 800", "below 850")) + self.assertIn('return "destination"', result) + source_method = ast.parse(self.source).body[0].body[0].body[0] + result_method = ast.parse(result).body[0].body[0].body[0] + self.assertEqual( + ast.get_source_segment(self.source, source_method), + ast.get_source_segment(result, result_method), + ) + + def test_guard_mismatch_is_rejected(self): + with self.assertRaisesRegex(ValueError, "guards differ"): + copy_test( + self.source.replace("if int.__doc__", "if enabled"), + self.destination, + "Finder.example", + ) + + def test_missing_or_ambiguous_definition_is_rejected(self): + for source in ("class Other: pass\n", self.source + self.source): + with self.subTest(source=source): + with self.assertRaisesRegex(ValueError, "Expected one class"): + copy_test(source, self.destination, "Finder.example") + duplicate = self.source.replace( + " def untouched(self):", + " def example():\n pass\n\n def untouched(self):", + ) + with self.assertRaisesRegex(ValueError, "Expected one method"): + copy_test(duplicate, self.destination, "Finder.example") + + def test_decorator_and_signature_changes_are_rejected(self): + decorated = self.source.replace( + " def example():", " @decorator\n def example():" + ) + with self.assertRaisesRegex(ValueError, "decorators differ"): + copy_test(decorated, self.destination, "Finder.example") + with self.assertRaisesRegex(ValueError, "signatures differ"): + copy_test( + self.source.replace("def example():", "def example(arg):"), + self.destination, + "Finder.example", + ) + + def test_nested_local_scope_is_not_selected(self): + local = "def factory():\n" + "".join( + " " + line for line in self.source.splitlines(True) + ) + with self.assertRaisesRegex(ValueError, "Expected one class"): + copy_test(local, self.destination, "Finder.example") + + def test_decorator_comments_must_match_exactly(self): + for source_marker, destination_marker in ( + ( + "@unittest.expectedFailure", + "@unittest.expectedFailure # TODO: RUSTPYTHON; reason", + ), + ( + "@unittest.expectedFailure # original reason", + "@unittest.expectedFailure # changed reason", + ), + ( + "@unittest.expectedFailure\n @other", + "@unittest.expectedFailure\n # marker rationale\n @other", + ), + ( + "@unittest.expectedFailureIf(\n condition,\n 'reason',\n )", + "@unittest.expectedFailureIf(\n condition, # project limitation\n 'reason',\n )", + ), + ): + with self.subTest(destination_marker=destination_marker): + source = self.source.replace( + " def example():", + f" {source_marker}\n def example():", + ) + destination = self.destination.replace( + " def example():", + f" {destination_marker}\n def example():", + ) + with self.assertRaisesRegex(ValueError, "decorator text differs"): + copy_test(source, destination, "Finder.example") + + def test_identical_decorator_comments_are_preserved(self): + prefix = ( + " @unittest.expectedFailure # TODO: RUSTPYTHON; reason\n" + " # Keep this rationale with the marker.\n" + " def example():" + ) + for declaration in ("def example():", "async def example():"): + with self.subTest(declaration=declaration): + source = self.source.replace(" def example():", prefix) + destination = self.destination.replace(" def example():", prefix) + source = source.replace("def example():", declaration) + destination = destination.replace("def example():", declaration) + self.assertEqual( + copy_test(source, destination, "Finder.example"), + destination.replace("below 800", "below 850"), + ) + + def test_cli_does_not_write_on_decorator_comment_mismatch(self): + with tempfile.TemporaryDirectory() as temp: + source = Path(temp) / "source.py" + destination = Path(temp) / "destination.py" + source.write_text( + self.source.replace( + " def example():", + " @unittest.expectedFailure\n def example():", + ) + ) + destination.write_text( + self.destination.replace( + " def example():", + " @unittest.expectedFailure # TODO: RUSTPYTHON; reason\n" + " def example():", + ) + ) + before = destination.read_bytes() + with ( + contextlib.redirect_stderr(io.StringIO()) as errors, + self.assertRaises(SystemExit), + ): + main( + [ + str(source), + str(destination), + "--test", + "Finder.example", + "--source-sha256", + hashlib.sha256(source.read_bytes()).hexdigest(), + "--destination-sha256", + hashlib.sha256(before).hexdigest(), + ] + ) + self.assertIn("decorator text differs", errors.getvalue()) + self.assertEqual(destination.read_bytes(), before) + + def test_return_annotation_and_type_parameters_are_rejected(self): + for declaration in ("def example() -> int:", "def example[T]():"): + with self.subTest(declaration=declaration): + source = self.source.replace("def example():", declaration) + with self.assertRaisesRegex(ValueError, "signatures differ"): + copy_test(source, self.destination, "Finder.example") + + def test_cli_hashes_and_no_write_on_stale_input(self): + with tempfile.TemporaryDirectory() as temp: + source = Path(temp) / "source.py" + destination = Path(temp) / "destination.py" + source.write_bytes(self.source.encode()) + destination.write_bytes(self.destination.encode()) + args = [ + str(source), + str(destination), + "--test", + "Finder.example", + "--source-sha256", + hashlib.sha256(source.read_bytes()).hexdigest(), + "--destination-sha256", + "0" * 64, + ] + with ( + contextlib.redirect_stderr(io.StringIO()), + self.assertRaises(SystemExit), + ): + main(args) + self.assertEqual(destination.read_bytes(), self.destination.encode()) + args[-1] = hashlib.sha256(destination.read_bytes()).hexdigest() + with contextlib.redirect_stdout(io.StringIO()) as output: + self.assertEqual(main(args), 0) + self.assertIn('"source_sha256"', output.getvalue()) + self.assertEqual( + destination.read_text(), + self.destination.replace("below 800", "below 850"), + ) diff --git a/scripts/update_lib/tests/test_patch_spec.py b/scripts/update_lib/tests/test_patch_spec.py index 92a0a35a153..2c6bb99ff84 100644 --- a/scripts/update_lib/tests/test_patch_spec.py +++ b/scripts/update_lib/tests/test_patch_spec.py @@ -67,6 +67,38 @@ async def test_async(self): class TestExtractPatches(unittest.TestCase): """Tests for extract_patches function.""" + def test_class_conditional_method_markers_are_preserved(self): + code = """class Owner: + if enabled: + @unittest.expectedFailure # TODO: RUSTPYTHON + def test_value(self): + return 42 +""" + patches = extract_patches(code) + self.assertIn("test_value", patches["Owner"]) + clean = code.replace( + " @unittest.expectedFailure # TODO: RUSTPYTHON\n", "" + ) + result = apply_patches(clean, patches) + self.assertIn(" @unittest.expectedFailure", result) + method = list(iter_tests(ast.parse(result)))[0][1] + self.assertEqual(method.col_offset, 8) + self.assertEqual(ast.dump(method.body[0]), "Return(value=Constant(value=42))") + + def test_nested_class_markers_do_not_target_module_names(self): + code = """class Container: + class Target: + @unittest.expectedFailure # TODO: RUSTPYTHON + def test_value(self): + pass +def factory(): + class Target: + @unittest.expectedFailure # TODO: RUSTPYTHON + def test_value(self): + pass +""" + self.assertEqual(extract_patches(code), {}) + def test_extract_expected_failure(self): """Test extracting @unittest.expectedFailure decorator.""" code = f""" @@ -177,6 +209,106 @@ def test_one(self): class TestApplyPatches(unittest.TestCase): """Tests for apply_patches function.""" + def test_conditional_class_existing_method(self): + code = """import unittest +if enabled: + class TestConditional: + def test_value(self): + assert 1 + 1 == 2 +""" + result = apply_patches( + code, + { + "TestConditional": { + "test_value": [ + PatchSpec( + UtMethod.ExpectedFailureIf, "condition", "known limitation" + ) + ] + } + }, + ) + original = list(iter_tests(ast.parse(code)))[0][1] + updated = list(iter_tests(ast.parse(result)))[0][1] + self.assertEqual(ast.dump(original.body[0]), ast.dump(updated.body[0])) + self.assertEqual(updated.col_offset, 8) + self.assertIn(" @unittest.expectedFailureIf", result) + + def test_conditional_class_inherited_override(self): + code = """import unittest +class Parent: + async def test_value(self): + return 42 +if enabled: + class Child(Parent): + pass +""" + result = apply_patches( + code, {"Child": {"test_value": [PatchSpec(UtMethod.ExpectedSuccess)]}} + ) + methods = list(iter_tests(ast.parse(result))) + self.assertEqual([cls.name for cls, _ in methods], ["Parent", "Child"]) + self.assertIsInstance(methods[1][1], ast.AsyncFunctionDef) + self.assertIn(" return await super().test_value()", result) + self.assertIn(" @unittest.expectedSuccess", result) + + def test_duplicate_conditional_class_names_are_rejected(self): + code = """if enabled: + class Duplicate: + pass +else: + class Duplicate: + pass +""" + with self.assertRaisesRegex(ValueError, "Ambiguous.*Duplicate"): + apply_patches( + code, + {"Duplicate": {"test_value": [PatchSpec(UtMethod.ExpectedFailure)]}}, + ) + + def test_local_class_names_do_not_shadow_module_class(self): + code = """import unittest +class Target: + def test_value(self): + return 1 +class Container: + class Target: + def test_value(self): + return 2 +def factory(): + class Target: + def test_value(self): + return 3 +""" + result = apply_patches( + code, {"Target": {"test_value": [PatchSpec(UtMethod.ExpectedFailure)]}} + ) + tree = ast.parse(result) + targets = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.ClassDef) and node.name == "Target" + ] + self.assertEqual( + [len(node.body[0].decorator_list) for node in targets], [1, 0, 0] + ) + + def test_duplicate_conditional_method_names_are_rejected(self): + code = """class Owner: + if enabled: + def test_value(self): + return 1 + else: + def test_value(self): + return 2 +""" + with self.assertRaisesRegex( + ValueError, "Ambiguous method name: Owner.test_value" + ): + apply_patches( + code, {"Owner": {"test_value": [PatchSpec(UtMethod.ExpectedFailure)]}} + ) + def test_apply_expected_failure(self): """Test applying @unittest.expectedFailure.""" code = """import unittest diff --git a/src/shell/helper.rs b/src/shell/helper.rs index 15fc75563dc..1b6d3d72af8 100644 --- a/src/shell/helper.rs +++ b/src/shell/helper.rs @@ -1,14 +1,14 @@ #![cfg_attr(target_arch = "wasm32", allow(dead_code))] use rustpython_vm::{ AsObject, PyResult, TryFromObject, VirtualMachine, - builtins::{PyDictRef, PyStrRef}, + builtins::{PyAnyDictRef, PyStrRef}, function::ArgIterable, identifier, }; pub(super) struct ShellHelper<'vm> { vm: &'vm VirtualMachine, - globals: PyDictRef, + globals: PyAnyDictRef, } fn reverse_string(s: &mut String) { @@ -54,7 +54,7 @@ fn split_idents_on_dot(line: &str) -> Option<(usize, Vec)> { } impl<'vm> ShellHelper<'vm> { - pub(super) const fn new(vm: &'vm VirtualMachine, globals: PyDictRef) -> Self { + pub(super) const fn new(vm: &'vm VirtualMachine, globals: PyAnyDictRef) -> Self { ShellHelper { vm, globals } }