diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py index e76169c69..b7580867e 100644 --- a/Lib/test/test_listcomps.py +++ b/Lib/test/test_listcomps.py @@ -691,7 +691,6 @@ class ListComprehensionTest(unittest.TestCase): """ self._check_in_scopes(code, {"value": [1, None]}) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_frame_locals(self): code = """ val = "a" in [sys._getframe().f_locals for a in [0]][0] diff --git a/crates/capi/src/ceval.rs b/crates/capi/src/ceval.rs index 7de5b3cf6..c69701a86 100644 --- a/crates/capi/src/ceval.rs +++ b/crates/capi/src/ceval.rs @@ -126,8 +126,18 @@ pub extern "C" fn PyEval_GetLocals() -> *mut PyObject { let Some(frame) = vm.current_frame() else { return Ok(core::ptr::null_mut()); }; - let _ = frame.locals(vm)?; - Ok(frame.iframe().locals.as_object(vm).as_raw().cast_mut()) + let locals = frame.locals(vm)?; + if frame + .iframe() + .locals + .get() + .is_some_and(|mapping| mapping.obj().is(locals.obj())) + { + return Ok(locals.obj().as_raw().cast_mut()); + } + let cache = frame.locals_snapshot_cache(vm); + cache.merge_object(locals.into(), vm)?; + Ok(cache.as_object().as_raw().cast_mut()) }) } @@ -149,9 +159,49 @@ pub unsafe extern "C" fn PyEval_GetFuncDesc(func: *mut PyObject) -> *const c_cha #[cfg(test)] mod tests { + use alloc::ffi::CString; use pyo3::exceptions::PyException; use pyo3::prelude::*; + #[pyfunction] + fn legacy_local(name: &str) -> (usize, Option) { + let locals = super::PyEval_GetLocals(); + assert!(!locals.is_null()); + let name = CString::new(name).unwrap(); + let value = unsafe { crate::dictobject::PyDict_GetItemString(locals, name.as_ptr()) }; + let value = if value.is_null() { + None + } else { + Some(unsafe { crate::longobject::PyLong_AsLongLong(value) }) + }; + (locals as usize, value) + } + + #[pyfunction] + fn legacy_locals_is_frame_mapping() -> bool { + let locals = super::PyEval_GetLocals(); + rustpython_vm::vm::thread::with_current_vm(|vm| { + let frame = vm.current_frame().unwrap(); + core::ptr::eq( + locals.cast_const(), + frame.iframe().locals.as_object(vm).as_raw(), + ) + }) + } + + #[pyfunction] + fn frame_locals_are_fresh() -> bool { + let first = super::PyEval_GetFrameLocals(); + let second = super::PyEval_GetFrameLocals(); + assert!(!first.is_null() && !second.is_null()); + let different = first != second; + unsafe { + crate::refcount::_Py_DecRef(first); + crate::refcount::_Py_DecRef(second); + } + different + } + #[test] fn code_eval() { Python::attach(|py| { @@ -167,4 +217,55 @@ mod tests { assert!(err.is_instance_of::(py)); }) } + + #[test] + fn legacy_locals_cache_is_lazy_and_does_not_modify_namespaces() { + Python::attach(|py| { + let globals = pyo3::types::PyDict::new(py); + globals + .set_item("legacy_local", wrap_pyfunction!(legacy_local, py).unwrap()) + .unwrap(); + globals + .set_item( + "legacy_locals_is_frame_mapping", + wrap_pyfunction!(legacy_locals_is_frame_mapping, py).unwrap(), + ) + .unwrap(); + globals + .set_item( + "frame_locals_are_fresh", + wrap_pyfunction!(frame_locals_are_fresh, py).unwrap(), + ) + .unwrap(); + py.run( + c"\ +assert legacy_locals_is_frame_mapping() + +def optimized(): + x = 1 + first_ptr, first_x = legacy_local('x') + x = 2 + second_ptr, second_x = legacy_local('x') + return first_ptr == second_ptr, first_x, second_x, frame_locals_are_fresh() + +def optimized_with_custom_locals(): + x = 3 + legacy_local('x') + +optimized_result = optimized() +custom_locals = {} +exec(optimized_with_custom_locals.__code__, globals(), custom_locals) +hidden_result = [legacy_local('hidden_name')[1] for hidden_name in [7]][0] +hidden_leaked = 'hidden_name' in globals() +assert optimized_result == (True, 1, 2, True) +assert custom_locals == {} +assert hidden_result == 7 +assert not hidden_leaked +", + Some(&globals), + None, + ) + .unwrap(); + }) + } } diff --git a/crates/vm/src/builtins/code.rs b/crates/vm/src/builtins/code.rs index ef2720dd9..9739cfe8e 100644 --- a/crates/vm/src/builtins/code.rs +++ b/crates/vm/src/builtins/code.rs @@ -472,6 +472,10 @@ pub struct CoMonitoringData { #[pyclass(module = false, name = "code")] pub struct PyCode { pub code: CodeObject, + /// Slot-indexed names, equivalent to CPython's `co_localsplusnames`. + /// Derived once so frame-local proxy operations do not repeatedly scan + /// merged cell variables. + localsplus_names: Box<[&'static PyStrInterned]>, source_path: AtomicPtr, /// Version counter for lazy re-instrumentation. /// Compared against `PyGlobalState::instrumentation_version` at RESUME. @@ -494,9 +498,107 @@ impl Deref for PyCode { } } +fn build_localspluskinds( + varnames: &[&'static PyStrInterned], + cellvars: &[&'static PyStrInterned], + freevars: &[&'static PyStrInterned], + arg_counts: (u32, u32, u32), + flags: CodeFlags, + instructions: &CodeUnits, +) -> Result, usize> { + use rustpython_compiler_core::bytecode::{ + CO_FAST_ARG_KW, CO_FAST_ARG_POS, CO_FAST_ARG_VAR, CO_FAST_CELL, CO_FAST_FREE, + CO_FAST_HIDDEN, CO_FAST_LOCAL, OpArgState, + }; + + let num_merged_cells = cellvars + .iter() + .filter(|cell| varnames.iter().any(|local| *local == **cell)) + .count(); + let mut kinds = vec![0; varnames.len() + cellvars.len() - num_merged_cells + freevars.len()]; + + let (posonlyarg_count, arg_count, kwonlyarg_count) = arg_counts; + let positional_only = posonlyarg_count as usize; + let positional_or_keyword = arg_count.saturating_sub(posonlyarg_count) as usize; + let argument_kinds = [ + (positional_only, CO_FAST_ARG_POS), + (positional_or_keyword, CO_FAST_ARG_POS | CO_FAST_ARG_KW), + (kwonlyarg_count as usize, CO_FAST_ARG_KW), + ( + usize::from(flags.contains(CodeFlags::VARARGS)), + CO_FAST_ARG_VAR | CO_FAST_ARG_POS, + ), + ( + usize::from(flags.contains(CodeFlags::VARKEYWORDS)), + CO_FAST_ARG_VAR | CO_FAST_ARG_KW, + ), + (usize::MAX, 0), + ]; + let mut local_index = 0; + let mut argument_end = 0usize; + for (count, argument_kind) in argument_kinds { + argument_end = argument_end.saturating_add(count); + while local_index < argument_end && local_index < varnames.len() { + kinds[local_index] = CO_FAST_LOCAL | argument_kind; + local_index += 1; + } + } + + let mut dropped_cells = 0; + for (cell_index, cell) in cellvars.iter().enumerate() { + if let Some(local_index) = varnames.iter().position(|local| *local == *cell) { + kinds[local_index] |= CO_FAST_CELL; + dropped_cells += 1; + } else { + kinds[varnames.len() + cell_index - dropped_cells] = CO_FAST_CELL; + } + } + + let free_start = varnames.len() + cellvars.len() - num_merged_cells; + for kind in kinds.iter_mut().skip(free_start) { + *kind = CO_FAST_FREE; + } + + if !flags.contains(CodeFlags::OPTIMIZED) { + let mut arg_state = OpArgState::default(); + for unit in instructions.iter().copied() { + let (instruction, arg) = arg_state.get(unit); + if matches!(instruction, Instruction::LoadFastAndClear { .. }) { + let index = u32::from(arg) as usize; + let Some(kind) = kinds.get_mut(index) else { + return Err(index); + }; + *kind |= CO_FAST_HIDDEN; + } + } + } + + Ok(kinds.into_boxed_slice()) +} + impl PyCode { pub fn new(code: CodeObject) -> Self { let sp = code.source_path as *const PyStrInterned as *mut PyStrInterned; + let localsplus_names = { + let varname_ids = code + .varnames + .iter() + .map(|name| name.get_id()) + .collect::>(); + let names = code + .varnames + .iter() + .chain( + code.cellvars + .iter() + .filter(|name| !varname_ids.contains(&name.get_id())), + ) + .chain(code.freevars.iter()) + .copied() + .collect::>(); + debug_assert_eq!(names.len(), code.localspluskinds.len()); + names + }; // The only opcodes that call `vm.set_exception` (mutating the shared // exc_info slot); instrumented variants only replace these base opcodes // in place, so scanning the freshly-built stream is a sound predicate. @@ -512,6 +614,7 @@ impl PyCode { }); Self { code, + localsplus_names, source_path: AtomicPtr::new(sp), instrumentation_version: AtomicU64::new(0), monitoring_data: PyMutex::new(None), @@ -520,6 +623,14 @@ impl PyCode { } } + #[inline(always)] + pub(crate) fn localsplus_name(&self, index: usize) -> &'static PyStrInterned { + // Bytecode operands and frame-proxy indices are validated against + // localspluskinds, whose length is kept equal to localsplus_names by + // every CodeObject construction path. + self.localsplus_names[index] + } + pub fn source_path(&self) -> &'static PyStrInterned { // SAFETY: always points to a valid &'static PyStrInterned (interned strings are never deallocated) unsafe { &*self.source_path.load(Ordering::Relaxed) } @@ -805,43 +916,26 @@ impl Constructor for PyCode { )], > = vec![(loc, loc); instructions.len()].into_boxed_slice(); - // Build localspluskinds with cell-local merging - let localspluskinds = { - use rustpython_compiler_core::bytecode::*; - let nlocals = varnames.len(); - let ncells = cellvars.len(); - let nfrees = freevars.len(); - let numdropped = cellvars - .iter() - .filter(|cv| varnames.iter().any(|v| *v == **cv)) - .count(); - let nlocalsplus = nlocals + ncells - numdropped + nfrees; - let mut kinds = vec![0u8; nlocalsplus]; - for kind in kinds.iter_mut().take(nlocals) { - *kind = CO_FAST_LOCAL; - } - let mut cell_numdropped = 0usize; - for (i, cv) in cellvars.iter().enumerate() { - let merged_idx = varnames.iter().position(|v| **v == **cv); - if let Some(local_idx) = merged_idx { - kinds[local_idx] |= CO_FAST_CELL; - cell_numdropped += 1; - } else { - kinds[nlocals + i - cell_numdropped] = CO_FAST_CELL; - } - } - let free_start = nlocals + ncells - numdropped; - for i in 0..nfrees { - kinds[free_start + i] = CO_FAST_FREE; - } - kinds.into_boxed_slice() - }; + let flags = CodeFlags::from_bits_truncate(args.flags); + let localspluskinds = build_localspluskinds( + &varnames, + &cellvars, + &freevars, + (args.posonlyargcount, args.argcount, args.kwonlyargcount), + flags, + &instructions, + ) + .map_err(|index| { + vm.new_value_error(format!( + "code: LOAD_FAST_AND_CLEAR oparg {index} out of range" + )) + })?; // Build the CodeObject let code = CodeObject { instructions, locations, - flags: CodeFlags::from_bits_truncate(args.flags), + flags, posonlyarg_count: args.posonlyargcount, arg_count: args.argcount, kwonlyarg_count: args.kwonlyargcount, @@ -1347,12 +1441,12 @@ impl PyCode { .collect(), }; - let flags = match co_flags { + let flags = CodeFlags::from_bits_truncate(match co_flags { OptionalArg::Present(flags) => flags, OptionalArg::Missing => self.code.flags.bits(), - }; + }); - let varnames = match co_varnames { + let varname_objects = match co_varnames { OptionalArg::Present(varnames) => varnames, OptionalArg::Missing => self.code.varnames.iter().map(|s| s.to_object()).collect(), }; @@ -1390,6 +1484,8 @@ impl PyCode { .map(Vec::into_boxed_slice) }; + let varnames = intern_all(varname_objects, "co_varnames")?; + let cellvars = match co_cellvars { OptionalArg::Present(cellvars) => intern_all(cellvars, "co_cellvars")?, OptionalArg::Missing => self.code.cellvars.clone(), @@ -1400,10 +1496,11 @@ impl PyCode { OptionalArg::Missing => self.code.freevars.clone(), }; - // Validate co_nlocals if provided - if let OptionalArg::Present(nlocals) = co_nlocals - && nlocals as usize != varnames.len() - { + let nlocals = match co_nlocals { + OptionalArg::Present(nlocals) => nlocals as usize, + OptionalArg::Missing => self.code.varnames.len(), + }; + if nlocals != varnames.len() { return Err(vm.new_value_error(format!( "co_nlocals ({}) != len(co_varnames) ({})", nlocals, @@ -1411,6 +1508,20 @@ impl PyCode { ))); } + let localspluskinds = build_localspluskinds( + &varnames, + &cellvars, + &freevars, + (posonlyarg_count, arg_count, kwonlyarg_count), + flags, + &instructions, + ) + .map_err(|index| { + vm.new_value_error(format!( + "code: LOAD_FAST_AND_CLEAR oparg {index} out of range" + )) + })?; + // Handle linetable and exceptiontable let linetable = match co_linetable { OptionalArg::Present(linetable) => linetable.as_bytes().to_vec().into_boxed_slice(), @@ -1425,7 +1536,7 @@ impl PyCode { }; let new_code = CodeObject { - flags: CodeFlags::from_bits_truncate(flags), + flags, posonlyarg_count, arg_count, kwonlyarg_count, @@ -1441,10 +1552,10 @@ impl PyCode { locations: self.code.locations.clone(), constants: constants.into_iter().map(Literal).collect(), names: intern_all(names, "co_names")?, - varnames: intern_all(varnames, "co_varnames")?, + varnames, cellvars, freevars, - localspluskinds: self.code.localspluskinds.clone(), + localspluskinds, linetable, exceptiontable, }; diff --git a/crates/vm/src/builtins/dict.rs b/crates/vm/src/builtins/dict.rs index 014b8bba5..9837c7c86 100644 --- a/crates/vm/src/builtins/dict.rs +++ b/crates/vm/src/builtins/dict.rs @@ -269,7 +269,7 @@ impl PyDict { Ok(()) } - fn merge_dict( + pub(crate) fn merge_dict( &self, dict_other: PyDictRef, override_existing: bool, diff --git a/crates/vm/src/builtins/frame.rs b/crates/vm/src/builtins/frame.rs index e4b3aa2b1..3704e070c 100644 --- a/crates/vm/src/builtins/frame.rs +++ b/crates/vm/src/builtins/frame.rs @@ -764,6 +764,13 @@ impl Py { #[pymethod] // = frame_clear_impl fn clear(&self, vm: &VirtualMachine) -> PyResult<()> { + // Materialized stack frames are FrameObject-owned even while their + // source iframe is still executing. TLS lookup below only sees the + // calling thread, so attached_tid is the cross-thread-safe execution + // state check. + if self.iframe().attached_tid() != 0 { + return Err(vm.new_runtime_error("cannot clear an executing frame")); + } let owner = FrameOwner::from_i8( self.iframe() .owner @@ -791,44 +798,62 @@ impl Py { } } - // Clear fastlocals - // SAFETY: FrameObject is not executing (detached or stopped). - { - let fastlocals = unsafe { self.fastlocals_mut() }; - for slot in fastlocals.iter_mut() { - *slot = None; - } - } + // Move references out before dropping them. Their finalizers may + // re-enter this frame, so no locals borrow or cold-data lock may be + // held while they run. + let fastlocals = { + // SAFETY: FrameObject is not executing (detached or stopped). + let slots = unsafe { self.fastlocals_mut() }; + slots + .iter_mut() + .filter_map(Option::take) + .collect::>() + }; // Clear the evaluation stack and cell references self.clear_stack_and_cells(); - // Clear temporary refs - self.iframe().cold().temporary_refs.lock().clear(); - self.iframe().cold().f_locals_hidden_overlay.lock().take(); - self.iframe().cold().f_extra_locals.lock().take(); - self.iframe().cold().retained_back.lock().take(); + let cold = self.iframe().cold(); + let temporary_refs = { + let mut guard = cold.temporary_refs.lock(); + core::mem::take(&mut *guard) + }; + let extra_locals = { + let mut guard = cold.f_extra_locals.lock(); + guard.take() + }; + let locals_cache = { + let mut guard = cold.f_locals_cache.lock(); + guard.take() + }; + let overwritten = { + let mut guard = cold.f_overwritten_fast_locals.lock(); + core::mem::take(&mut *guard) + }; + let retained_back = { + let mut guard = cold.retained_back.lock(); + guard.take() + }; + drop(( + fastlocals, + temporary_refs, + extra_locals, + locals_cache, + overwritten, + retained_back, + )); Ok(()) } #[pygetset] fn f_locals(&self, vm: &VirtualMachine) -> PyResult { - // Optimized (function) frames expose a live write-through - // FrameLocalsProxy; class/module/exec frames expose their namespace - // mapping directly. - if self - .iframe() - .code() - .flags - .contains(bytecode::CodeFlags::OPTIMIZED) - { - self.check_locals_access(vm)?; + if self.uses_locals_proxy(vm)? { self.mark_escaped(); let proxy = crate::builtins::FrameLocalsProxy::new(self.to_owned()); Ok(proxy.into_ref(&vm.ctx).into()) } else { - self.f_locals_mapping(vm).map(Into::into) + Ok(self.iframe().locals.clone_mapping(vm).into()) } } diff --git a/crates/vm/src/builtins/frame_locals_proxy.rs b/crates/vm/src/builtins/frame_locals_proxy.rs index eaba0bca6..3fc150fd8 100644 --- a/crates/vm/src/builtins/frame_locals_proxy.rs +++ b/crates/vm/src/builtins/frame_locals_proxy.rs @@ -1,6 +1,7 @@ //! The `FrameLocalsProxy` type returned by `frame.f_locals` for optimized -//! (function) frames. Implements PEP 667 write-through semantics on top of the -//! frame's fast-local slots and an extra-locals side dict. +//! (function) frames and for unoptimized frames that currently have PEP 709 +//! hidden comprehension locals. Implements PEP 667 write-through semantics on +//! top of the frame's fast-local slots and an extra-locals side dict. use super::{PyDict, PyDictRef, PyType}; use crate::{ @@ -10,7 +11,7 @@ use crate::{ frame::FrameObjectRef, function::{FuncArgs, OptionalArg, PyArithmeticValue, PyComparisonValue}, object::{Traverse, TraverseFn}, - protocol::{PyMappingMethods, PyNumberMethods, PySequenceMethods}, + protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods}, recursion::ReprGuard, types::{ AsMapping, AsNumber, AsSequence, Comparable, Constructor, Iterable, PyComparisonOp, @@ -20,7 +21,12 @@ use crate::{ use rustpython_common::lock::LazyLock; use rustpython_common::wtf8::Wtf8Buf; -#[pyclass(module = false, name = "FrameLocalsProxy", traverse = "manual")] +#[pyclass( + module = false, + name = "FrameLocalsProxy", + unhashable = true, + traverse = "manual" +)] #[derive(Debug)] pub struct FrameLocalsProxy { frame: FrameObjectRef, @@ -48,8 +54,16 @@ impl FrameLocalsProxy { self.frame.framelocalsproxy_snapshot(vm) } + fn items_vec(&self, vm: &VirtualMachine) -> PyResult> { + self.frame.framelocalsproxy_items(vm) + } + fn keys_vec(&self, vm: &VirtualMachine) -> PyResult> { - Ok(self.snapshot(vm)?.into_iter().map(|(k, _)| k).collect()) + Ok(self + .items_vec(vm)? + .into_iter() + .map(|(key, _)| key) + .collect()) } } @@ -108,7 +122,7 @@ impl FrameLocalsProxy { } fn __len__(&self, vm: &VirtualMachine) -> PyResult { - Ok(self.snapshot(vm)?.__len__()) + Ok(self.items_vec(vm)?.len()) } #[pymethod] @@ -118,14 +132,18 @@ impl FrameLocalsProxy { #[pymethod] fn values(&self, vm: &VirtualMachine) -> PyResult { - let values = self.snapshot(vm)?.into_iter().map(|(_, v)| v).collect(); + let values = self + .items_vec(vm)? + .into_iter() + .map(|(_, value)| value) + .collect(); Ok(vm.ctx.new_list(values).into()) } #[pymethod] fn items(&self, vm: &VirtualMachine) -> PyResult { let items = self - .snapshot(vm)? + .items_vec(vm)? .into_iter() .map(|(k, v)| vm.ctx.new_tuple(vec![k, v]).into()) .collect(); @@ -175,17 +193,19 @@ impl FrameLocalsProxy { } fn update_from(&self, other: &PyObject, vm: &VirtualMachine) -> PyResult<()> { - let items: Vec<(PyObjectRef, PyObjectRef)> = - if let Some(dict) = other.downcast_ref::() { - dict.into_iter().collect() - } else if let Some(proxy) = other.downcast_ref::() { - proxy.snapshot(vm)?.into_iter().collect() - } else { - return Err( - vm.new_type_error("update() argument must be dict or another FrameLocalsProxy") - ); - }; - for (key, value) in items { + if other.downcast_ref::().is_none() && other.downcast_ref::().is_none() { + return Err( + vm.new_type_error("update() argument must be dict or another FrameLocalsProxy") + ); + } + // CPython deliberately uses the mapping protocol here, including + // overridden keys()/__getitem__ on dict subclasses. + let keys = other + .get_attr(vm.ctx.intern_str("keys"), vm)? + .call((), vm)? + .get_iter(vm)?; + while let PyIterReturn::Return(key) = keys.next(vm)? { + let value = other.get_item(&*key, vm)?; self.frame.framelocalsproxy_setitem(key, value, vm)?; } Ok(()) @@ -199,18 +219,35 @@ impl FrameLocalsProxy { } fn __ior__(zelf: PyRef, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { + if other.downcast_ref::().is_none() && other.downcast_ref::().is_none() { + return Ok(vm.ctx.not_implemented()); + } zelf.update_from(&other, vm)?; Ok(zelf.into()) } fn __or__(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { - let base = self.snapshot(vm)?; - vm._or(base.as_object(), &other) + if other.downcast_ref::().is_none() && other.downcast_ref::().is_none() { + return Ok(vm.ctx.not_implemented()); + } + let result = self.snapshot(vm)?; + if other.downcast_ref::().is_some() { + // PyDict_Update reads a dict subclass's stored entries directly; + // it does not dispatch to overridden mapping methods. + result.merge_dict(other.downcast().unwrap(), true, vm)?; + } else { + result.merge_object(other, vm)?; + } + Ok(result.into()) } fn __ror__(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { - let base = self.snapshot(vm)?; - vm._or(&other, base.as_object()) + let Some(other) = other.downcast_ref::() else { + return Ok(vm.ctx.not_implemented()); + }; + let result = other.copy().into_ref(&vm.ctx); + result.merge_object(self.snapshot(vm)?.into(), vm)?; + Ok(result.into()) } #[pymethod] @@ -297,12 +334,14 @@ impl Comparable for FrameLocalsProxy { vm: &VirtualMachine, ) -> PyResult { op.eq_only(|| { + if let Some(other) = other.downcast_ref::() { + return Ok(PyComparisonValue::Implemented(zelf.frame.is(&other.frame))); + } + if other.downcast_ref::().is_none() { + return Ok(PyComparisonValue::NotImplemented); + } let self_dict: PyObjectRef = zelf.snapshot(vm)?.into(); - let other_obj = match other.downcast_ref::() { - Some(proxy) => proxy.snapshot(vm)?.into(), - None => other.to_owned(), - }; - let res = self_dict.rich_compare(other_obj, PyComparisonOp::Eq, vm)?; + let res = self_dict.rich_compare(other.to_owned(), PyComparisonOp::Eq, vm)?; PyArithmeticValue::from_object(vm, res) .map(|o| o.try_to_bool(vm)) .transpose() diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 4cc4912ec..9715ded94 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -26,7 +26,7 @@ use crate::{ function::{ArgMapping, Either, FuncArgs, KwArgs, PyMethodFlags}, object::PyAtomicBorrow, object::{Traverse, TraverseFn}, - protocol::{PyIter, PyIterReturn, PyMapping}, + protocol::{PyIter, PyIterReturn}, scope::Scope, sliceable::SliceableSequenceOp, stdlib::{_typing, builtins, sys::monitoring}, @@ -619,6 +619,24 @@ impl LocalsPlus { } } + /// Move active stack references out as owned references without running + /// finalizers while this locals-plus storage is mutably borrowed. + fn take_stack_object_refs(&mut self) -> Vec { + let mut refs = Vec::with_capacity(self.stack_top as usize); + while self.stack_top > 0 { + if let Some(value) = self.stack_pop() { + refs.push(value.to_pyobj()); + } + } + refs + } + + /// Extract every owned Python reference for GC's deferred-drop phase. + fn clear_into(&mut self, out: &mut Vec) { + out.extend(self.take_stack_object_refs()); + out.extend(self.fastlocals_mut().iter_mut().filter_map(Option::take)); + } + /// Drain stack elements from `from` to the end, returning an iterator /// that yields `Option` in forward order and shrinks the stack. fn stack_drain( @@ -757,6 +775,10 @@ impl FrameLocals { pub fn as_object(&self, vm: &VirtualMachine) -> &PyObject { self.get_or_create(vm).obj() } + + fn take(&mut self) -> Option { + self.inner.take() + } } impl fmt::Debug for FrameLocals { @@ -793,8 +815,11 @@ pub(crate) struct FrameColdData { pub trace_lines: PyMutex, pub trace_opcodes: PyMutex, pub temporary_refs: PyMutex>, - pub f_locals_hidden_overlay: PyMutex>, pub f_extra_locals: PyMutex>, + /// Backing storage for the borrowed reference returned by the legacy + /// `PyEval_GetLocals()` C API. It is populated only by that adapter. + pub f_locals_cache: PyMutex>, + pub f_overwritten_fast_locals: PyMutex>, pub escaped: atomic::AtomicBool, pub retained_back: PyMutex>, pub pending_stack_pops: PyAtomic, @@ -814,8 +839,9 @@ impl Default for FrameColdData { trace_lines: PyMutex::new(true), trace_opcodes: PyMutex::new(false), temporary_refs: PyMutex::new(Vec::new()), - f_locals_hidden_overlay: PyMutex::new(None), f_extra_locals: PyMutex::new(None), + f_locals_cache: PyMutex::new(None), + f_overwritten_fast_locals: PyMutex::new(Vec::new()), escaped: atomic::AtomicBool::new(false), retained_back: PyMutex::new(None), pending_stack_pops: Default::default(), @@ -1523,21 +1549,52 @@ unsafe impl Traverse for FrameObject { if let Some(cold) = iframe.cold_opt() { cold.trace.traverse(tracer_fn); cold.temporary_refs.traverse(tracer_fn); - cold.f_locals_hidden_overlay.traverse(tracer_fn); cold.f_extra_locals.traverse(tracer_fn); + cold.f_locals_cache.traverse(tracer_fn); + cold.f_overwritten_fast_locals.traverse(tracer_fn); cold.retained_back.traverse(tracer_fn); } } - fn clear(&mut self, _out: &mut Vec) { - // Drop the interpreter frame and owned reference anchors so GC - // cycle collection can reclaim the referenced objects. The payload - // is left as a trivially-droppable husk for the freelist. - drop(self.iframe.get_mut().take()); - self.owned_code.take(); - self.owned_globals.take(); - self.owned_builtins.take(); - self.owned_func_obj.take(); + fn clear(&mut self, out: &mut Vec) { + // Extract every child before dropping the frame husk. GC drops `out` + // after this exclusive payload borrow ends, so re-entrant finalizers + // cannot alias frame storage or run under one of its locks. + if let Some(mut iframe) = self.iframe.get_mut().take() { + iframe.localsplus.clear_into(out); + if let Some(locals) = iframe.locals.take() { + out.push(locals.into()); + } + if let Some(cold) = iframe.cold.take() { + let cold = *cold; + if let Some(trace) = cold.trace.into_inner() { + out.push(trace); + } + out.extend(cold.temporary_refs.into_inner()); + if let Some(extra) = cold.f_extra_locals.into_inner() { + out.push(extra.into()); + } + if let Some(cache) = cold.f_locals_cache.into_inner() { + out.push(cache.into()); + } + out.extend(cold.f_overwritten_fast_locals.into_inner()); + if let Some(back) = cold.retained_back.into_inner() { + out.push(back.into()); + } + } + } + if let Some(code) = self.owned_code.take() { + out.push(code.into()); + } + if let Some(globals) = self.owned_globals.take() { + out.push(globals.into()); + } + if let Some(builtins) = self.owned_builtins.take() { + out.push(builtins); + } + if let Some(func) = self.owned_func_obj.take() { + out.push(func); + } } } @@ -1694,22 +1751,39 @@ impl FrameObject { pub(crate) fn clear_stack_and_cells(&self) { // SAFETY: Called when frame is not executing (generator closed). // Cell refs in fastlocals[nlocals..] are cleared by clear_locals_and_stack(). - unsafe { - self.iframe_mut().localsplus.stack_clear(); - } + let refs = unsafe { self.iframe_mut().localsplus.take_stack_object_refs() }; + drop(refs); } /// Clear locals and stack after generator/coroutine close. /// Releases references held by the frame, matching _PyFrame_ClearLocals. pub(crate) fn clear_locals_and_stack(&self) { self.clear_stack_and_cells(); - // SAFETY: FrameObject is not executing (generator closed). - let fastlocals = unsafe { self.iframe_mut().localsplus.fastlocals_mut() }; - for slot in fastlocals.iter_mut() { - *slot = None; - } - self.iframe().cold().f_locals_hidden_overlay.lock().take(); - self.iframe().cold().f_extra_locals.lock().take(); + // Move references out before dropping them. Their finalizers may + // re-enter this frame, so no locals borrow or cold-data lock may be + // held while they run. + let fastlocals = { + // SAFETY: FrameObject is not executing (generator closed). + let slots = unsafe { self.iframe_mut().localsplus.fastlocals_mut() }; + slots + .iter_mut() + .filter_map(Option::take) + .collect::>() + }; + let cold = self.iframe().cold(); + let extra_locals = { + let mut guard = cold.f_extra_locals.lock(); + guard.take() + }; + let locals_cache = { + let mut guard = cold.f_locals_cache.lock(); + guard.take() + }; + let overwritten = { + let mut guard = cold.f_overwritten_fast_locals.lock(); + core::mem::take(&mut *guard) + }; + drop((fastlocals, extra_locals, locals_cache, overwritten)); } /// Store a borrowed back-reference to the owning generator/coroutine. @@ -1820,135 +1894,55 @@ impl FrameObject { } } - fn has_active_hidden_locals(&self) -> bool { + /// True when this frame currently has bound PEP 709 hidden comprehension + /// locals. Matches CPython `_PyFrame_HasHiddenLocals`. + pub(crate) fn has_active_hidden_locals(&self) -> bool { use rustpython_compiler_core::bytecode::{CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN}; let code = self.iframe().code(); - // SAFETY: reached from `locals()` on the thread running this frame. + // SAFETY: callers first pass through `check_locals_access`, or this is + // `locals()` on the thread running this frame. let fastlocals = unsafe { self.live_fastlocals() }; - let is_optimized = code.flags.contains(bytecode::CodeFlags::OPTIMIZED); - !is_optimized - && code.localspluskinds.iter().enumerate().any(|(i, &kind)| { - if kind & CO_FAST_HIDDEN == 0 { - return false; - } - match fastlocals[i].as_ref() { - None => false, - Some(obj) => { - if kind & (CO_FAST_CELL | CO_FAST_FREE) != 0 { - obj.downcast_ref::() - .is_none_or(|cell| cell.get().is_some()) - } else { - true - } + code.localspluskinds.iter().enumerate().any(|(i, &kind)| { + if kind & CO_FAST_HIDDEN == 0 { + return false; + } + match fastlocals[i].as_ref() { + None => false, + Some(obj) => { + if kind & (CO_FAST_CELL | CO_FAST_FREE) != 0 { + obj.downcast_ref::() + .is_none_or(|cell| cell.get().is_some()) + } else { + true } } - }) + } + }) } - fn sync_visible_locals_to_mapping( - &self, - locals_map: PyMapping<'_>, - vm: &VirtualMachine, - ) -> PyResult<()> { - use rustpython_compiler_core::bytecode::{ - CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, - }; - // SAFETY: Either the frame is not executing (caller checked owner), - // or we're in a trace callback on the same thread that's executing. - let code = self.iframe().code(); - let fastlocals = unsafe { self.live_fastlocals() }; + #[inline] + pub(crate) fn has_hidden_local_slots(&self) -> bool { + use rustpython_compiler_core::bytecode::CO_FAST_HIDDEN; + self.iframe() + .code() + .localspluskinds + .iter() + .any(|kind| kind & CO_FAST_HIDDEN != 0) + } - // Iterate through all localsplus slots using localspluskinds - let nlocalsplus = code.localspluskinds.len(); - let nfrees = code.freevars.len(); - let free_start = nlocalsplus - nfrees; - let is_optimized = code.flags.contains(bytecode::CodeFlags::OPTIMIZED); - - // Track which non-merged cellvar index we're at - let mut nonmerged_cell_idx = 0; - - for (i, &kind) in code.localspluskinds.iter().enumerate() { - if kind & CO_FAST_HIDDEN != 0 { - // Hidden variables are only skipped when their slot is empty. - // After a comprehension restores values, they should appear in locals(). - let slot_empty = match fastlocals[i].as_ref() { - None => true, - Some(obj) => { - if kind & (CO_FAST_CELL | CO_FAST_FREE) != 0 { - // If it's a PyCell, check if the cell is empty. - // If it's a raw value (merged cell during inlined comp), not empty. - obj.downcast_ref::() - .is_some_and(|cell| cell.get().is_none()) - } else { - false - } - } - }; - if slot_empty { - continue; - } - } - - // CPython only syncs fastlocals/cells into locals() for function - // scope; class/module scope just returns the namespace dict as-is - // (_PyFrame_GetLocals, Objects/frameobject.c). - if !is_optimized && kind & (CO_FAST_CELL | CO_FAST_FREE) != 0 { - continue; - } - - // Get the name for this slot - let name = if kind & CO_FAST_LOCAL != 0 { - code.varnames[i] - } else if kind & CO_FAST_FREE != 0 { - code.freevars[i - free_start] - } else if kind & CO_FAST_CELL != 0 { - // Non-merged cell: find the name by skipping merged cellvars - let mut found_name = None; - let mut skip = nonmerged_cell_idx; - for cv in &code.cellvars { - let is_merged = code.varnames.contains(cv); - if !is_merged { - if skip == 0 { - found_name = Some(*cv); - break; - } - skip -= 1; - } - } - nonmerged_cell_idx += 1; - match found_name { - Some(n) => n, - None => continue, - } - } else { - continue; - }; - - // Get the value - let value = if kind & (CO_FAST_CELL | CO_FAST_FREE) != 0 { - // Cell or free var: extract value from PyCell. - // During inlined comprehensions, a merged cell slot may hold a raw - // value (not a PyCell) after LOAD_FAST_AND_CLEAR + STORE_FAST. - fastlocals[i].as_ref().and_then(|obj| { - if let Some(cell) = obj.downcast_ref::() { - cell.get() - } else { - Some(obj.clone()) - } - }) - } else { - // Regular local - fastlocals[i].clone() - }; - - let result = locals_map.ass_subscript(name, value, vm); - match result { - Ok(()) => {} - Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => {} - Err(e) => return Err(e), - } + /// Decide whether Python-visible locals use a write-through proxy rather + /// than the frame's namespace mapping. + pub(crate) fn uses_locals_proxy(&self, vm: &VirtualMachine) -> PyResult { + let is_optimized = self + .iframe() + .code() + .flags + .contains(bytecode::CodeFlags::OPTIMIZED); + if !is_optimized && !self.has_hidden_local_slots() { + return Ok(false); } - Ok(()) + self.check_locals_access(vm)?; + Ok(is_optimized || self.has_active_hidden_locals()) } /// Reject locals access for a frame that is executing on another thread. @@ -1995,55 +1989,21 @@ impl FrameObject { )) } - pub fn f_locals_mapping(&self, vm: &VirtualMachine) -> PyResult { - self.check_locals_access(vm)?; - if !self.has_active_hidden_locals() { - self.iframe().cold().f_locals_hidden_overlay.lock().take(); - return self.locals(vm); - } - - let overlay_dict = { - let mut overlay = self.iframe().cold().f_locals_hidden_overlay.lock(); - match overlay.as_ref() { - Some(dict) => dict.clone(), - None => { - let dict = vm.ctx.new_dict(); - *overlay = Some(dict.clone()); - dict - } - } - }; - PyDict::clear(&overlay_dict); - let overlay = ArgMapping::from_dict_exact(overlay_dict.clone()); - self.sync_visible_locals_to_mapping(overlay.mapping(), vm)?; - Ok(ArgMapping::from_dict_exact(overlay_dict)) - } - pub fn locals(&self, vm: &VirtualMachine) -> PyResult { - let mapping = if self.has_active_hidden_locals() { - // Match CPython's locals() behavior for frames with PEP 709 hidden - // locals: return a fresh snapshot instead of the backing mapping. - let overlay = ArgMapping::from_dict_exact(vm.ctx.new_dict()); - self.sync_visible_locals_to_mapping(overlay.mapping(), vm)?; - overlay + if self.uses_locals_proxy(vm)? { + Ok(ArgMapping::from_dict_exact( + self.framelocalsproxy_snapshot(vm)?, + )) } else { - self.sync_visible_locals_to_mapping(self.iframe().locals.mapping(vm), vm)?; - self.iframe().locals.clone_mapping(vm) - }; - self.fold_extra_locals(&mapping, vm)?; - Ok(mapping) + Ok(self.iframe().locals.clone_mapping(vm)) + } } - /// Copy the frame's extra-locals side storage (proxy keys that are not - /// fast locals) into `mapping`. No-op when nothing was ever stored. - fn fold_extra_locals(&self, mapping: &ArgMapping, vm: &VirtualMachine) -> PyResult<()> { - let extra = self.iframe().cold().f_extra_locals.lock().clone(); - if let Some(extra) = extra { - for (key, value) in &extra { - mapping.mapping().ass_subscript(&key, Some(value), vm)?; - } - } - Ok(()) + /// Return the lazily allocated dict used by APIs that must keep a + /// frame-owned locals snapshot. Ordinary VM execution never accesses it. + pub fn locals_snapshot_cache(&self, vm: &VirtualMachine) -> PyDictRef { + let mut cache = self.iframe().cold().f_locals_cache.lock(); + cache.get_or_insert_with(|| vm.ctx.new_dict()).clone() } /// Read a fast-local slot's visible value, dereferencing cells. `None` if @@ -2083,23 +2043,45 @@ impl FrameObject { .get(i) .copied() .unwrap_or(0); - // SAFETY: callers first pass through `check_locals_access`. - // Use live source iframe if available so writes reach the - // executing frame's actual local variables. let live = self.find_live_source_iframe(); - let fastlocals = if !live.is_null() { - unsafe { &mut *live.cast_mut() }.localsplus.fastlocals_mut() + let old_value = if !live.is_null() { + // SAFETY: callers first pass through `check_locals_access`. + unsafe { (*live).localsplus.fastlocals()[i].clone() } } else { - unsafe { self.iframe_mut().localsplus.fastlocals_mut() } + // SAFETY: callers first pass through `check_locals_access`. + unsafe { self.iframe_ref().localsplus.fastlocals()[i].clone() } }; + if kind & (CO_FAST_CELL | CO_FAST_FREE) != 0 - && let Some(obj) = fastlocals[i].as_ref() - && let Some(cell) = obj.downcast_ref::() + && let Some(cell) = old_value + .as_ref() + .and_then(|obj| obj.clone().downcast::().ok()) { cell.set(Some(value)); return; } - fastlocals[i] = Some(value); + + if old_value.as_ref().is_some_and(|old| old.is(&value)) { + return; + } + + // Keep the overwritten value alive with the frame. Its finalizer may + // re-enter this frame, so it must not run while `fastlocals` is + // mutably borrowed. CPython uses f_overwritten_fast_locals likewise. + let old_value = if !live.is_null() { + // SAFETY: callers first pass through `check_locals_access`. + unsafe { &mut *live.cast_mut() }.localsplus.fastlocals_mut()[i].replace(value) + } else { + // SAFETY: callers first pass through `check_locals_access`. + unsafe { self.iframe_mut().localsplus.fastlocals_mut()[i].replace(value) } + }; + if let Some(old_value) = old_value { + self.iframe() + .cold() + .f_overwritten_fast_locals + .lock() + .push(old_value); + } } /// Resolve `key` to a fast-local slot index, or `None` if it names no fast @@ -2113,10 +2095,13 @@ impl FrameObject { vm: &VirtualMachine, ) -> PyResult> { use rustpython_compiler_core::bytecode::CO_FAST_HIDDEN; - // The proxy hashes the key first; an unhashable key raises TypeError. - key.hash(vm)?; + // Hash first both for hashability and to match dict-key equivalence. + let key_hash = key.hash(vm)?; for (i, &kind) in self.iframe().code().localspluskinds.iter().enumerate() { let name = localsplus_name(self.iframe().code(), i); + if name.as_object().hash(vm)? != key_hash { + continue; + } if !name .as_object() .rich_compare_bool(key, PyComparisonOp::Eq, vm)? @@ -2134,14 +2119,45 @@ impl FrameObject { Ok(None) } - /// Build a fresh ordered snapshot dict of the proxy's visible contents: - /// bound fast locals in localsplus order followed by extra locals. + /// Return the proxy's items in CPython iteration order. Fast locals come + /// first, followed by extra locals. The two groups may contain equal keys. + pub(crate) fn framelocalsproxy_items( + &self, + vm: &VirtualMachine, + ) -> PyResult> { + self.check_locals_access(vm)?; + let code = self.iframe().code(); + let mut items = Vec::with_capacity(code.localspluskinds.len()); + for i in 0..code.localspluskinds.len() { + if let Some(value) = self.framelocalsproxy_getval(i) { + let name = localsplus_name(code, i).to_owned().into(); + items.push((name, value)); + } + } + let extra = self.iframe().cold().f_extra_locals.lock().clone(); + if let Some(extra) = extra { + items.extend(&extra); + } + Ok(items) + } + + /// Build a dict as `dict(FrameLocalsProxy)` does in CPython. Insert fast + /// locals first, then only extra keys that are still missing, so a + /// colliding fast local keeps precedence. pub(crate) fn framelocalsproxy_snapshot(&self, vm: &VirtualMachine) -> PyResult { self.check_locals_access(vm)?; let dict = vm.ctx.new_dict(); - let mapping = ArgMapping::from_dict_exact(dict.clone()); - self.sync_visible_locals_to_mapping(mapping.mapping(), vm)?; - self.fold_extra_locals(&mapping, vm)?; + let code = self.iframe().code(); + for i in 0..code.localspluskinds.len() { + if let Some(value) = self.framelocalsproxy_getval(i) { + let key: PyObjectRef = localsplus_name(code, i).to_owned().into(); + dict.set_item(&*key, value, vm)?; + } + } + let extra = self.iframe().cold().f_extra_locals.lock().clone(); + if let Some(extra) = extra { + dict.merge_object_if_missing(extra.into(), vm)?; + } Ok(dict) } @@ -2607,43 +2623,7 @@ fn specialization_nonnegative_compact_index(i: &PyInt, vm: &VirtualMachine) -> O /// Get the variable name for a localsplus index of `code`. fn localsplus_name(code: &PyCode, idx: usize) -> &'static PyStrInterned { - use rustpython_compiler_core::bytecode::{CO_FAST_CELL, CO_FAST_FREE, CO_FAST_LOCAL}; - let nlocals = code.varnames.len(); - let kind = code.localspluskinds.get(idx).copied().unwrap_or(0); - if kind & CO_FAST_LOCAL != 0 { - // Merged cell or regular local: name is in varnames - code.varnames[idx] - } else if kind & CO_FAST_FREE != 0 { - // Free var: slots are at the end of localsplus - let nlocalsplus = code.localspluskinds.len(); - let nfrees = code.freevars.len(); - let free_start = nlocalsplus - nfrees; - code.freevars[idx - free_start] - } else if kind & CO_FAST_CELL != 0 { - // Non-merged cell: count how many non-merged cell slots are before - // this index to find the corresponding cellvars entry. - // Non-merged cellvars appear in their original order (skipping merged ones). - let nonmerged_pos = code.localspluskinds[nlocals..idx] - .iter() - .filter(|&&k| k == CO_FAST_CELL) - .count(); - // Skip merged cellvars to find the right one - let mut cv_idx = 0; - let mut nonmerged_count = 0; - for (i, name) in code.cellvars.iter().enumerate() { - let is_merged = code.varnames.contains(name); - if !is_merged { - if nonmerged_count == nonmerged_pos { - cv_idx = i; - break; - } - nonmerged_count += 1; - } - } - code.cellvars[cv_idx] - } else { - code.varnames[idx] - } + code.localsplus_name(idx) } /// Free a finished call frame's data stack storage. diff --git a/extra_tests/snippets/vm_frame_locals.py b/extra_tests/snippets/vm_frame_locals.py new file mode 100644 index 000000000..608167c2a --- /dev/null +++ b/extra_tests/snippets/vm_frame_locals.py @@ -0,0 +1,301 @@ +"""Tests for frame-local mappings and PEP 667 snapshot semantics.""" + +import sys +import threading +import unittest + + +class FrameLocalsProxyTest(unittest.TestCase): + def test_extra_locals_do_not_leak_into_unoptimized_namespace(self): + namespace = {"sys": sys} + exec( + """ +x = 1 +proxy = [sys._getframe().f_locals for a in [0]][0] +proxy["x"] = 2 +frame_value = sys._getframe().f_locals["x"] +namespace_value = x +proxy_value = proxy["x"] +""", + namespace, + ) + + self.assertEqual(namespace["frame_value"], 1) + self.assertEqual(namespace["namespace_value"], 1) + self.assertEqual(namespace["proxy_value"], 2) + + def test_hidden_locals_take_precedence_over_extra_locals(self): + namespace = {"sys": sys} + exec( + """ +value = [ + ( + sys._getframe().f_locals.__setitem__("a", 99), + locals()["a"], + )[1] + for a in [7] +][0] +""", + namespace, + ) + + self.assertEqual(namespace["value"], 7) + + def test_proxy_views_preserve_colliding_extra_local(self): + def inspect_proxy(proxy): + proxy["a"] = 99 + return ( + len(proxy), + list(proxy.keys()), + list(proxy.values()), + list(proxy.items()), + dict(proxy), + repr(proxy), + proxy["a"], + ) + + namespace = {"inspect_proxy": inspect_proxy, "sys": sys} + exec( + "observed = [inspect_proxy(sys._getframe().f_locals) for a in [7]][0]", + namespace, + ) + + self.assertEqual( + namespace["observed"], + ( + 2, + ["a", "a"], + [7, 99], + [("a", 7), ("a", 99)], + {"a": 7}, + "{'a': 7}", + 7, + ), + ) + + def test_optimized_locals_returns_fresh_snapshots(self): + def snapshots(): + before = locals() + x = 1 + after = locals() + return before, after, before is after + + before, after, same = snapshots() + self.assertEqual(before, {}) + self.assertEqual(after, {"before": before, "x": 1}) + self.assertFalse(same) + + def test_proxy_equality_uses_frame_identity(self): + def make_proxy(): + return sys._getframe().f_locals + + first = make_proxy() + second = make_proxy() + self.assertNotEqual(first, second) + + def same_frame(): + frame = sys._getframe() + return frame.f_locals == frame.f_locals + + self.assertTrue(same_frame()) + + def test_proxy_keys_use_dict_hash_and_equality_rules(self): + class EqualWithDifferentHash: + def __hash__(self): + return hash("x") ^ 1 + + def __eq__(self, other): + return other == "x" + + def exercise(): + x = 1 + proxy = sys._getframe().f_locals + key = EqualWithDifferentHash() + proxy[key] = 2 + return x, proxy[key], proxy["x"] + + self.assertEqual(exercise(), (1, 2, 1)) + + def test_proxy_hash_comparison_and_union_dispatch(self): + class Reflected: + def __ror__(self, other): + return ("reflected", type(other).__name__) + + def __eq__(self, other): + return ("equal", type(other).__name__) + + def exercise(): + x = 1 + proxy = sys._getframe().f_locals + reflected = Reflected() + inplace = proxy + inplace |= reflected + proxy |= {"x": 3, "extra": 4} + return ( + type(proxy).__hash__, + proxy | {"right": 2}, + {"left": 2} | proxy, + proxy | reflected, + inplace, + proxy == reflected, + x, + proxy["extra"], + ) + + ( + hash_method, + merged, + reflected_merged, + reflected_or, + reflected_inplace, + reflected_equal, + x, + extra, + ) = exercise() + self.assertIsNone(hash_method) + self.assertEqual(merged["x"], 3) + self.assertEqual(merged["right"], 2) + self.assertEqual(reflected_merged["left"], 2) + self.assertEqual(reflected_merged["x"], 3) + self.assertEqual(reflected_or, ("reflected", "FrameLocalsProxy")) + self.assertEqual(reflected_inplace, ("reflected", "FrameLocalsProxy")) + self.assertEqual(reflected_equal, ("equal", "FrameLocalsProxy")) + self.assertEqual((x, extra), (3, 4)) + + def get_proxy(): + return sys._getframe().f_locals + + proxy = get_proxy() + with self.assertRaises(TypeError): + hash(proxy) + with self.assertRaisesRegex(TypeError, "FrameLocalsProxy.*list"): + proxy | [] + + def test_dict_subclass_protocol_matches_cpython(self): + class DictSubclass(dict): + def keys(self): + return ["virtual"] + + def __getitem__(self, key): + if key == "virtual": + return 42 + return super().__getitem__(key) + + def exercise(): + proxy = sys._getframe().f_locals + proxy.update(DictSubclass(real=1)) + merged = proxy | DictSubclass(real=2) + reflected = DictSubclass(real=3) | proxy + proxy |= DictSubclass(real=4) + return ( + proxy["virtual"], + "real" in proxy, + merged["real"], + reflected["real"], + ) + + self.assertEqual(exercise(), (42, False, 2, 3)) + + def test_clear_releases_values_without_finalizer_reentrancy_deadlock(self): + events = [] + + class ReentrantFinalizer: + def __init__(self, frame, key): + self.frame = frame + self.key = key + + def __del__(self): + events.append(self.frame.f_locals.get(self.key)) + + def frame_with_extra(): + frame = sys._getframe() + frame.f_locals["extra"] = ReentrantFinalizer(frame, "extra") + return frame + + frame = frame_with_extra() + frame.clear() + self.assertEqual(events, [None]) + + events.clear() + + def frame_with_overwritten_fast_local(): + frame = sys._getframe() + # Occupies a fast-local slot so the write below overwrites it. + value = ReentrantFinalizer(frame, "value") # noqa: F841 + frame.f_locals["value"] = 42 + return frame + + frame = frame_with_overwritten_fast_local() + self.assertEqual(events, []) + frame.clear() + self.assertEqual(len(events), 1) + + def test_unoptimized_frame_locals_are_readable_from_another_thread(self): + ready = threading.Event() + release = threading.Event() + namespace = {"ready": ready, "release": release, "sys": sys} + + thread = threading.Thread( + target=exec, + args=( + "frame = sys._getframe()\nready.set()\nrelease.wait()", + namespace, + ), + ) + thread.start() + self.assertTrue(ready.wait(5)) + try: + self.assertIs(namespace["frame"].f_locals, namespace) + finally: + release.set() + thread.join(5) + self.assertFalse(thread.is_alive()) + + def test_clear_executing_frame_from_another_thread(self): + ready = threading.Event() + release = threading.Event() + frames = [] + + def worker(): + frames.append(sys._getframe()) + ready.set() + release.wait() + + thread = threading.Thread(target=worker) + thread.start() + self.assertTrue(ready.wait(5)) + try: + with self.assertRaisesRegex(RuntimeError, "executing frame"): + frames[0].clear() + finally: + release.set() + thread.join(5) + self.assertFalse(thread.is_alive()) + + +class CodeReplaceLocalsPlusTest(unittest.TestCase): + def test_rejects_varnames_that_disagree_with_default_nlocals(self): + def function(value): + return value + + code = function.__code__ + with self.assertRaises(ValueError): + code.replace(co_varnames=()) + with self.assertRaises(ValueError): + code.replace(co_varnames=("value", "extra")) + + def test_rebuilds_metadata_for_replaced_cellvars_and_freevars(self): + def outer(value): + return lambda: value + + code = outer.__code__ + replaced_cells = code.replace(co_cellvars=("replacement",)) + self.assertEqual(replaced_cells.co_cellvars, ("replacement",)) + + plain_code = (lambda value: value).__code__ + replaced_frees = plain_code.replace(co_freevars=("replacement",)) + self.assertEqual(replaced_frees.co_freevars, ("replacement",)) + + +if __name__ == "__main__": + unittest.main()