diff --git a/crates/vm/src/builtins/mod.rs b/crates/vm/src/builtins/mod.rs index c7404fcd1..a938b9b3d 100644 --- a/crates/vm/src/builtins/mod.rs +++ b/crates/vm/src/builtins/mod.rs @@ -75,6 +75,8 @@ pub(crate) mod type_; pub use type_::{PyType, PyTypeRef}; pub(crate) mod range; pub use range::PyRange; +pub(crate) mod sentinel; +pub use sentinel::PySentinel; pub(crate) mod set; pub use set::{PyFrozenSet, PySet}; pub(crate) mod singletons; diff --git a/crates/vm/src/builtins/sentinel.rs b/crates/vm/src/builtins/sentinel.rs new file mode 100644 index 000000000..eefc33a58 --- /dev/null +++ b/crates/vm/src/builtins/sentinel.rs @@ -0,0 +1,118 @@ +use super::{PyStrRef, PyType, union_}; +use crate::{ + Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, + class::PyClassImpl, + object::{PyAtomicRef, Traverse, TraverseFn}, + protocol::PyNumberMethods, + types::{AsNumber, Constructor, Representable}, +}; + +/// Create a unique sentinel object with the given name. +#[pyclass(module = false, name = "sentinel", traverse = "manual")] +#[derive(Debug)] +pub struct PySentinel { + #[pymember(name = "__name__", type = "object_ex")] + name: PyStrRef, + #[pymember(name = "__module__", type = "object_ex", writable)] + module: PyAtomicRef>, + repr: Option, +} + +// SAFETY: Visit each owned Python reference once, without cloning it. +unsafe impl Traverse for PySentinel { + fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) { + self.name.traverse(tracer_fn); + self.module.traverse(tracer_fn); + self.repr.traverse(tracer_fn); + } + + fn clear(&mut self, out: &mut Vec) { + // String subclasses may own references back to this sentinel. + let name = core::mem::replace(&mut self.name, Context::genesis().empty_str.to_owned()); + out.push(name.into()); + if let Some(module) = self.module.store(None) { + out.push(module); + } + if let Some(repr) = self.repr.take() { + out.push(repr.into()); + } + } +} + +impl PyPayload for PySentinel { + #[inline] + fn class(ctx: &Context) -> &'static Py { + ctx.types.sentinel_type + } +} + +#[derive(FromArgs)] +pub struct SentinelArgs { + #[pyarg(positional, error_msg = "sentinel() argument 1 must be str")] + name: PyStrRef, + #[pyarg( + named, + optional, + error_msg = "sentinel() argument 'repr' must be str or None" + )] + repr: Option, +} + +impl Constructor for PySentinel { + type Args = SentinelArgs; + + fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { + let module = if let Some(frame) = vm.current_frame() + && let Some(function) = frame.iframe().func_obj() + { + // Use the module captured by the caller function, including the + // synthetic functions used for module, eval, and exec code. + function.get_attr(identifier!(vm, __module__), vm)? + } else { + vm.ctx.none() + }; + Ok(Self { + name: args.name, + module: PyAtomicRef::from(Some(module)), + repr: args.repr, + }) + } +} + +#[pyclass(with(Constructor, Representable, AsNumber), flags(HAVE_GC))] +impl PySentinel { + #[pymethod] + fn __copy__(zelf: PyRef) -> PyRef { + zelf + } + + #[pymethod] + fn __deepcopy__(zelf: PyRef, _memo: PyObjectRef) -> PyRef { + zelf + } + + #[pymethod] + fn __reduce__(zelf: &Py) -> PyStrRef { + zelf.name.clone() + } +} + +impl Representable for PySentinel { + fn repr(zelf: &Py, _vm: &VirtualMachine) -> PyResult { + Ok(zelf.repr.as_ref().unwrap_or(&zelf.name).clone()) + } +} + +impl AsNumber for PySentinel { + fn as_number() -> &'static PyNumberMethods { + static AS_NUMBER: PyNumberMethods = PyNumberMethods { + or: Some(|a, b, vm| union_::or_op(a.to_owned(), b.to_owned(), vm)), + ..PyNumberMethods::NOT_IMPLEMENTED + }; + &AS_NUMBER + } +} + +pub(crate) fn init(context: &'static Context) { + PySentinel::extend_class(context, context.types.sentinel_type); +} diff --git a/crates/vm/src/builtins/union.rs b/crates/vm/src/builtins/union.rs index 83ea3bb54..42d21d144 100644 --- a/crates/vm/src/builtins/union.rs +++ b/crates/vm/src/builtins/union.rs @@ -202,6 +202,7 @@ impl Py { fn is_unionable(obj: &PyObject, vm: &VirtualMachine) -> bool { let cls = obj.class(); cls.is(vm.ctx.types.none_type) + || cls.is(vm.ctx.types.sentinel_type) || obj.downcastable::() || cls.fast_issubclass(vm.ctx.types.generic_alias_type) || cls.is(vm.ctx.types.union_type) diff --git a/crates/vm/src/stdlib/builtins.rs b/crates/vm/src/stdlib/builtins.rs index fbced8423..d899f1ffc 100644 --- a/crates/vm/src/stdlib/builtins.rs +++ b/crates/vm/src/stdlib/builtins.rs @@ -1479,6 +1479,7 @@ pub fn init_module(vm: &VirtualMachine, module: &Py) { "object" => ctx.types.object_type.to_owned(), "property" => ctx.types.property_type.to_owned(), "range" => ctx.types.range_type.to_owned(), + "sentinel" => ctx.types.sentinel_type.to_owned(), "set" => ctx.types.set_type.to_owned(), "slice" => ctx.types.slice_type.to_owned(), "staticmethod" => ctx.types.staticmethod_type.to_owned(), diff --git a/crates/vm/src/types/slot.rs b/crates/vm/src/types/slot.rs index 03eb703f5..50d3a7a73 100644 --- a/crates/vm/src/types/slot.rs +++ b/crates/vm/src/types/slot.rs @@ -253,6 +253,7 @@ bitflagset::bitflagset! { const IMMUTABLETYPE = 8; const HEAPTYPE = 9; const BASETYPE = 10; + const HAVE_GC = 14; const METHOD_DESCRIPTOR = 17; const IS_ABSTRACT = 20; // Built-in types that match the subject itself in pattern matching diff --git a/crates/vm/src/types/zoo.rs b/crates/vm/src/types/zoo.rs index 0d3c5d0fb..9672e4cf3 100644 --- a/crates/vm/src/types/zoo.rs +++ b/crates/vm/src/types/zoo.rs @@ -4,8 +4,8 @@ use crate::{ asyncgenerator, bool_, builtin_func, bytearray, bytes, capsule, classmethod, code, complex, coroutine, descriptor, dict, enumerate, filter, float, frame, frame_locals_proxy, function, generator, genericalias, getset, int, interpolation, iter, list, map, mappingproxy, memory, - module, namespace, object, property, pystr, range, set, singletons, slice, staticmethod, - super_, template, traceback, tuple, + module, namespace, object, property, pystr, range, sentinel, set, singletons, slice, + staticmethod, super_, template, traceback, tuple, type_::{self, PyType}, union_, weakproxy, weakref, zip, }, @@ -64,6 +64,7 @@ pub struct TypeZoo { pub memoryviewiterator_type: &'static Py, pub tuple_type: &'static Py, pub tuple_iterator_type: &'static Py, + pub sentinel_type: &'static Py, pub set_type: &'static Py, pub set_iterator_type: &'static Py, pub staticmethod_type: &'static Py, @@ -145,6 +146,7 @@ impl TypeZoo { memoryview_type: memory::PyMemoryView::init_builtin_type(), property_type: property::PyProperty::init_builtin_type(), range_type: range::PyRange::init_builtin_type(), + sentinel_type: sentinel::PySentinel::init_builtin_type(), set_type: set::PySet::init_builtin_type(), slice_type: slice::PySlice::init_builtin_type(), staticmethod_type: staticmethod::PyStaticMethod::init_builtin_type(), @@ -265,6 +267,7 @@ impl TypeZoo { weakref::init(context); weakproxy::init(context); singletons::init(context); + sentinel::init(context); module::init(context); namespace::init(context); mappingproxy::init(context);