PEP 649 annotation phase 2

crates/codegen/src/compile.rs

@@ -3384,6 +3384,94 @@ impl Compiler {
         Ok(num_annotations)
     }
 
+    /// Compile function annotations as a closure (PEP 649)
+    /// Returns true if an __annotate__ closure was created
+    /// NOTE: This requires symbol table support for annotation scopes.
+    /// Currently unused - kept for future implementation reference.
+    #[allow(dead_code, clippy::cast_possible_truncation)]
+    fn compile_annotations_closure(
+        &mut self,
+        func_name: &str,
+        parameters: &Parameters,
+        returns: Option<&Expr>,
+    ) -> CompileResult<bool> {
+        // Count annotations first
+        let parameters_iter = core::iter::empty()
+            .chain(&parameters.posonlyargs)
+            .chain(&parameters.args)
+            .chain(&parameters.kwonlyargs)
+            .map(|x| &x.parameter)
+            .chain(parameters.vararg.as_deref())
+            .chain(parameters.kwarg.as_deref());
+
+        let num_annotations: u32 = parameters_iter.filter(|p| p.annotation.is_some()).count()
+            as u32
+            + if returns.is_some() { 1 } else { 0 };
+
+        if num_annotations == 0 {
+            return Ok(false);
+        }
+
+        // Create a new scope for the __annotate__ function
+        let annotate_name = format!("<annotate of {func_name}>");
+        self.push_output(
+            bytecode::CodeFlags::OPTIMIZED | bytecode::CodeFlags::NEWLOCALS,
+            0, // posonlyarg_count
+            1, // arg_count (format parameter)
+            0, // kwonlyarg_count
+            annotate_name,
+        )?;
+
+        // Add 'format' parameter to varnames
+        self.current_code_info()
+            .metadata
+            .varnames
+            .insert("format".to_owned());
+
+        // Compile annotations inside the new scope
+        let parameters_iter = core::iter::empty()
+            .chain(&parameters.posonlyargs)
+            .chain(&parameters.args)
+            .chain(&parameters.kwonlyargs)
+            .map(|x| &x.parameter)
+            .chain(parameters.vararg.as_deref())
+            .chain(parameters.kwarg.as_deref());
+
+        for param in parameters_iter {
+            if let Some(annotation) = &param.annotation {
+                self.emit_load_const(ConstantData::Str {
+                    value: self.mangle(param.name.as_str()).into_owned().into(),
+                });
+                self.compile_annotation(annotation)?;
+            }
+        }
+
+        // Handle return annotation
+        if let Some(annotation) = returns {
+            self.emit_load_const(ConstantData::Str {
+                value: "return".into(),
+            });
+            self.compile_annotation(annotation)?;
+        }
+
+        // Build the map and return it
+        emit!(
+            self,
+            Instruction::BuildMap {
+                size: num_annotations,
+            }
+        );
+        emit!(self, Instruction::ReturnValue);
+
+        // Exit the scope and get the code object
+        let annotate_code = self.exit_scope();
+
+        // Make a closure from the code object
+        self.make_closure(annotate_code, bytecode::MakeFunctionFlags::empty())?;
+
+        Ok(true)
+    }
+
     // = compiler_function
     #[allow(clippy::too_many_arguments)]
     fn compile_function_def(
@@ -3449,6 +3537,9 @@ impl Compiler {
         }
 
         // Compile annotations
+        // TODO: Full PEP 649 deferred annotation compilation requires symbol table changes.
+        // Currently using immediate evaluation (like PEP 563 without string conversion).
+        // The __annotate__ infrastructure is in place in function.rs, module.rs, type.rs.
         let mut annotations_flag = bytecode::MakeFunctionFlags::empty();
         let num_annotations = self.visit_annotations(parameters, returns)?;
         if num_annotations > 0 {

crates/compiler-core/src/bytecode/oparg.rs

@@ -327,6 +327,8 @@ bitflags! {
         const KW_ONLY_DEFAULTS = 0x04;
         const DEFAULTS = 0x08;
         const TYPE_PARAMS = 0x10;
+        /// PEP 649: __annotate__ function closure (instead of __annotations__ dict)
+        const ANNOTATE = 0x20;
     }
 }
 

crates/vm/src/builtins/function.rs

@@ -401,6 +401,12 @@ impl PyFunction {
                 ))
             })?;
             *self.type_params.lock() = type_params;
+        } else if attr == bytecode::MakeFunctionFlags::ANNOTATE {
+            // PEP 649: Store the __annotate__ function closure
+            if !attr_value.is_callable() {
+                return Err(vm.new_type_error("__annotate__ must be callable".to_owned()));
+            }
+            *self.annotate.lock() = Some(attr_value);
         } else {
             unreachable!("This is a compiler bug");
         }
@@ -588,36 +594,50 @@ impl PyFunction {
 
     #[pygetset]
     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);
-                }
+        // First check if we have cached annotations
+        {
+            let annotations = self.annotations.lock();
+            if let Some(ref ann) = *annotations {
+                return Ok(ann.clone());
             }
-            // 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())
+        // Check for callable __annotate__ and clone it before calling
+        let annotate_fn = {
+            let annotate = self.annotate.lock();
+            if let Some(ref func) = *annotate {
+                if func.is_callable() {
+                    Some(func.clone())
+                } else {
+                    None
+                }
+            } else {
+                None
+            }
+        };
+
+        // Release locks before calling __annotate__ to avoid deadlock
+        if let Some(annotate_fn) = annotate_fn {
+            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()
+                    ))
+                })?;
+
+            // Cache the result
+            *self.annotations.lock() = Some(ann_dict.clone());
+            return Ok(ann_dict);
+        }
+
+        // No __annotate__ or not callable, create empty dict
+        let new_dict = vm.ctx.new_dict();
+        *self.annotations.lock() = Some(new_dict.clone());
+        Ok(new_dict)
     }
 
     #[pygetset(setter)]

crates/vm/src/builtins/type.rs

@@ -899,10 +899,8 @@ impl PyType {
 
         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__));
-        }
+        // Always clear cached annotations when __annotate__ is updated
+        attrs.swap_remove(identifier!(vm, __annotations_cache__));
 
         Ok(())
     }