PEP 649 annotation phase 1

crates/vm/src/builtins/function.rs

@@ -36,7 +36,8 @@ pub struct PyFunction {
     name: PyMutex<PyStrRef>,
     qualname: PyMutex<PyStrRef>,
     type_params: PyMutex<PyTupleRef>,
-    annotations: PyMutex<PyDictRef>,
+    annotations: PyMutex<Option<PyDictRef>>,
+    annotate: PyMutex<Option<PyObjectRef>>,
     module: PyMutex<PyObjectRef>,
     doc: PyMutex<PyObjectRef>,
     #[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<PyDictRef> {
+        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::<crate::builtins::PyDict>()
+                            .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::<crate::builtins::PyDict>().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]

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<Self>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
+        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<Self>,
+        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<Self>, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
+        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<Self>,
+        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 {

crates/vm/src/builtins/type.rs

@@ -851,6 +851,62 @@ impl PyType {
         Ok(())
     }
 
+    #[pygetset]
+    fn __annotate__(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
+        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<PyObjectRef>, 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<PyObjectRef> {
         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(())

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__,