From 51b628642cc04a1a073a8e7de5c7f4a2274476d3 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Tue, 13 Jan 2026 23:03:57 +0900 Subject: [PATCH] PEP 649 annotation phase 1 --- crates/vm/src/builtins/function.rs | 93 ++++++++++++++++++++-- crates/vm/src/builtins/module.rs | 97 ++++++++++++++++++++++- crates/vm/src/builtins/type.rs | 123 +++++++++++++++++++++++------ crates/vm/src/vm/context.rs | 3 + 4 files changed, 285 insertions(+), 31 deletions(-) diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index 27a3b8fee..01ef987ae 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -36,7 +36,8 @@ pub struct PyFunction { name: PyMutex, qualname: PyMutex, type_params: PyMutex, - annotations: PyMutex, + annotations: PyMutex>, + annotate: PyMutex>, module: PyMutex, doc: PyMutex, #[cfg(feature = "jit")] @@ -98,7 +99,8 @@ impl PyFunction { name, qualname: PyMutex::new(qualname), type_params: PyMutex::new(vm.ctx.empty_tuple.clone()), - annotations: PyMutex::new(vm.ctx.new_dict()), + annotations: PyMutex::new(None), + annotate: PyMutex::new(None), module: PyMutex::new(module), doc: PyMutex::new(doc), #[cfg(feature = "jit")] @@ -375,7 +377,7 @@ impl PyFunction { ))); } }; - *self.annotations.lock() = annotations; + *self.annotations.lock() = Some(annotations); } else if attr == bytecode::MakeFunctionFlags::CLOSURE { // For closure, we need special handling // The closure tuple contains cell objects @@ -585,13 +587,90 @@ impl PyFunction { } #[pygetset] - fn __annotations__(&self) -> PyDictRef { - self.annotations.lock().clone() + fn __annotations__(&self, vm: &VirtualMachine) -> PyResult { + let mut annotations = self.annotations.lock(); + let annotate = self.annotate.lock(); + + if annotations.is_none() { + // If we have a callable __annotate__, call it to get annotations + if let Some(ref annotate_fn) = *annotate { + if annotate_fn.is_callable() { + // Call __annotate__(1) where 1 = Format.VALUE + let one = vm.ctx.new_int(1); + let ann_dict = annotate_fn.call((one,), vm)?; + let ann_dict = + ann_dict + .downcast::() + .map_err(|obj| { + vm.new_type_error(format!( + "__annotate__ returned non-dict of type '{}'", + obj.class().name() + )) + })?; + *annotations = Some(ann_dict.clone()); + return Ok(ann_dict); + } + } + // No __annotate__ or not callable, create empty dict + let new_dict = vm.ctx.new_dict(); + *annotations = Some(new_dict.clone()); + return Ok(new_dict); + } + + Ok(annotations.clone().unwrap()) } #[pygetset(setter)] - fn set___annotations__(&self, annotations: PyDictRef) { - *self.annotations.lock() = annotations + fn set___annotations__(&self, value: PySetterValue, vm: &VirtualMachine) -> PyResult<()> { + match value { + PySetterValue::Assign(value) => { + if vm.is_none(&value) { + *self.annotations.lock() = None; + } else { + let annotations = + value.downcast::().map_err(|_| { + vm.new_type_error("__annotations__ must be set to a dict object") + })?; + *self.annotations.lock() = Some(annotations); + } + // Clear __annotate__ when __annotations__ is set + *self.annotate.lock() = None; + } + PySetterValue::Delete => { + *self.annotations.lock() = None; + *self.annotate.lock() = None; + } + } + Ok(()) + } + + #[pygetset] + fn __annotate__(&self, vm: &VirtualMachine) -> PyObjectRef { + self.annotate + .lock() + .clone() + .unwrap_or_else(|| vm.ctx.none()) + } + + #[pygetset(setter)] + fn set___annotate__(&self, value: PySetterValue, vm: &VirtualMachine) -> PyResult<()> { + match value { + PySetterValue::Assign(value) => { + if vm.is_none(&value) { + *self.annotate.lock() = Some(value); + } else if value.is_callable() { + *self.annotate.lock() = Some(value); + // Clear cached __annotations__ when __annotate__ is set + *self.annotations.lock() = None; + } else { + return Err(vm.new_type_error("__annotate__ must be callable or None")); + } + } + PySetterValue::Delete => { + return Err(vm.new_type_error("__annotate__ cannot be deleted")); + } + } + Ok(()) } #[pygetset] diff --git a/crates/vm/src/builtins/module.rs b/crates/vm/src/builtins/module.rs index faa6e4813..2f0592f4e 100644 --- a/crates/vm/src/builtins/module.rs +++ b/crates/vm/src/builtins/module.rs @@ -4,7 +4,7 @@ use crate::{ builtins::{PyStrInterned, pystr::AsPyStr}, class::PyClassImpl, convert::ToPyObject, - function::{FuncArgs, PyMethodDef}, + function::{FuncArgs, PyMethodDef, PySetterValue}, types::{GetAttr, Initializer, Representable}, }; @@ -182,6 +182,101 @@ impl PyModule { let attrs = dict.into_iter().map(|(k, _v)| k).collect(); Ok(attrs) } + + #[pygetset] + fn __annotate__(zelf: &Py, vm: &VirtualMachine) -> PyResult { + let dict = zelf.dict(); + // Get __annotate__ from dict, set to None if not present + if let Some(annotate) = dict.get_item_opt(identifier!(vm, __annotate__), vm)? { + Ok(annotate) + } else { + let none = vm.ctx.none(); + dict.set_item(identifier!(vm, __annotate__), none.clone(), vm)?; + Ok(none) + } + } + + #[pygetset(setter)] + fn set___annotate__( + zelf: &Py, + value: PySetterValue, + vm: &VirtualMachine, + ) -> PyResult<()> { + match value { + PySetterValue::Assign(value) => { + if !vm.is_none(&value) && !value.is_callable() { + return Err(vm.new_type_error("__annotate__ must be callable or None")); + } + let dict = zelf.dict(); + dict.set_item(identifier!(vm, __annotate__), value.clone(), vm)?; + // Clear __annotations__ if value is not None + if !vm.is_none(&value) { + dict.del_item(identifier!(vm, __annotations__), vm).ok(); + } + Ok(()) + } + PySetterValue::Delete => Err(vm.new_type_error("cannot delete __annotate__ attribute")), + } + } + + #[pygetset] + fn __annotations__(zelf: &Py, vm: &VirtualMachine) -> PyResult { + let dict = zelf.dict(); + + // Check if __annotations__ is already in dict + if let Some(annotations) = dict.get_item_opt(identifier!(vm, __annotations__), vm)? { + return Ok(annotations); + } + + // Get __annotate__ and call it if callable + let annotations = + if let Some(annotate) = dict.get_item_opt(identifier!(vm, __annotate__), vm)? { + if annotate.is_callable() { + // Call __annotate__(1) where 1 is FORMAT_VALUE + let result = annotate.call((1i32,), vm)?; + if !result.class().is(vm.ctx.types.dict_type) { + return Err(vm.new_type_error(format!( + "__annotate__ returned non-dict of type '{}'", + result.class().name() + ))); + } + result + } else { + vm.ctx.new_dict().into() + } + } else { + vm.ctx.new_dict().into() + }; + + // Cache the result (TODO: check if module is initializing) + dict.set_item(identifier!(vm, __annotations__), annotations.clone(), vm)?; + Ok(annotations) + } + + #[pygetset(setter)] + fn set___annotations__( + zelf: &Py, + value: PySetterValue, + vm: &VirtualMachine, + ) -> PyResult<()> { + let dict = zelf.dict(); + match value { + PySetterValue::Assign(value) => { + dict.set_item(identifier!(vm, __annotations__), value, vm)?; + // Clear __annotate__ from dict + dict.del_item(identifier!(vm, __annotate__), vm).ok(); + Ok(()) + } + PySetterValue::Delete => { + if dict.del_item(identifier!(vm, __annotations__), vm).is_err() { + return Err(vm.new_attribute_error("__annotations__".to_owned())); + } + // Also clear __annotate__ + dict.del_item(identifier!(vm, __annotate__), vm).ok(); + Ok(()) + } + } + } } impl Initializer for PyModule { diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs index fed9af976..4bc326b38 100644 --- a/crates/vm/src/builtins/type.rs +++ b/crates/vm/src/builtins/type.rs @@ -851,6 +851,62 @@ impl PyType { Ok(()) } + #[pygetset] + fn __annotate__(&self, vm: &VirtualMachine) -> PyResult { + if !self.slots.flags.has_feature(PyTypeFlags::HEAPTYPE) { + return Err(vm.new_attribute_error(format!( + "type object '{}' has no attribute '__annotate__'", + self.name() + ))); + } + + let attrs = self.attributes.read(); + // First try __annotate__, in case that's been set explicitly + if let Some(annotate) = attrs.get(identifier!(vm, __annotate__)).cloned() { + return Ok(annotate); + } + // Then try __annotate_func__ + if let Some(annotate) = attrs.get(identifier!(vm, __annotate_func__)).cloned() { + return Ok(annotate); + } + drop(attrs); + + // Set None if not found + let none = vm.ctx.none(); + self.attributes + .write() + .insert(identifier!(vm, __annotate_func__), none.clone()); + Ok(none) + } + + #[pygetset(setter)] + fn set___annotate__(&self, value: Option, vm: &VirtualMachine) -> PyResult<()> { + if value.is_none() { + return Err(vm.new_type_error("cannot delete __annotate__ attribute".to_owned())); + } + let value = value.unwrap(); + + if self.slots.flags.has_feature(PyTypeFlags::IMMUTABLETYPE) { + return Err(vm.new_type_error(format!( + "cannot set '__annotate__' attribute of immutable type '{}'", + self.name() + ))); + } + + if !vm.is_none(&value) && !value.is_callable() { + return Err(vm.new_type_error("__annotate__ must be callable or None".to_owned())); + } + + let mut attrs = self.attributes.write(); + attrs.insert(identifier!(vm, __annotate_func__), value.clone()); + // Clear cached annotations if value is not None + if !vm.is_none(&value) { + attrs.swap_remove(identifier!(vm, __annotations_cache__)); + } + + Ok(()) + } + #[pygetset] fn __annotations__(&self, vm: &VirtualMachine) -> PyResult { if !self.slots.flags.has_feature(PyTypeFlags::HEAPTYPE) { @@ -860,20 +916,37 @@ impl PyType { ))); } - let __annotations__ = identifier!(vm, __annotations__); - let annotations = self.attributes.read().get(__annotations__).cloned(); + // First try __annotations__ (e.g. for "from __future__ import annotations") + let attrs = self.attributes.read(); + if let Some(annotations) = attrs.get(identifier!(vm, __annotations__)).cloned() { + return Ok(annotations); + } + // Then try __annotations_cache__ + if let Some(annotations) = attrs.get(identifier!(vm, __annotations_cache__)).cloned() { + return Ok(annotations); + } + drop(attrs); - let annotations = if let Some(annotations) = annotations { - annotations + // Get __annotate__ and call it if callable + let annotate = self.__annotate__(vm)?; + let annotations = if annotate.is_callable() { + // Call __annotate__(1) where 1 is FORMAT_VALUE + let result = annotate.call((1i32,), vm)?; + if !result.class().is(vm.ctx.types.dict_type) { + return Err(vm.new_type_error(format!( + "__annotate__ returned non-dict of type '{}'", + result.class().name() + ))); + } + result } else { - let annotations: PyObjectRef = vm.ctx.new_dict().into(); - let removed = self - .attributes - .write() - .insert(__annotations__, annotations.clone()); - debug_assert!(removed.is_none()); - annotations + vm.ctx.new_dict().into() }; + + // Cache the result in __annotations_cache__ + self.attributes + .write() + .insert(identifier!(vm, __annotations_cache__), annotations.clone()); Ok(annotations) } @@ -886,20 +959,24 @@ impl PyType { ))); } - let __annotations__ = identifier!(vm, __annotations__); + let mut attrs = self.attributes.write(); if let Some(value) = value { - self.attributes.write().insert(__annotations__, value); + attrs.insert(identifier!(vm, __annotations__), value); + // Clear __annotate__ when __annotations__ is set + attrs.swap_remove(identifier!(vm, __annotate__)); } else { - self.attributes - .read() - .get(__annotations__) - .cloned() - .ok_or_else(|| { - vm.new_attribute_error(format!( - "'{}' object has no attribute '__annotations__'", - self.name() - )) - })?; + // Delete + if attrs + .swap_remove(identifier!(vm, __annotations__)) + .is_none() + { + return Err(vm.new_attribute_error(format!( + "'{}' object has no attribute '__annotations__'", + self.name() + ))); + } + // Also clear __annotate__ + attrs.swap_remove(identifier!(vm, __annotate__)); } Ok(()) diff --git a/crates/vm/src/vm/context.rs b/crates/vm/src/vm/context.rs index 65c742e49..9978554bc 100644 --- a/crates/vm/src/vm/context.rs +++ b/crates/vm/src/vm/context.rs @@ -91,7 +91,10 @@ declare_const_name! { __all__, __and__, __anext__, + __annotate__, + __annotate_func__, __annotations__, + __annotations_cache__, __args__, __await__, __bases__,