Add get/set methods for function module/annotations (#5392)

---------

Co-authored-by: Jeong YunWon <jeong@youknowone.org>

vm/src/builtins/function.rs

@@ -38,6 +38,9 @@ pub struct PyFunction {
     type_params: PyMutex<PyTupleRef>,
     #[cfg(feature = "jit")]
     jitted_code: OnceCell<CompiledCode>,
+    annotations: PyMutex<PyDictRef>,
+    module: PyMutex<PyObjectRef>,
+    doc: PyMutex<PyObjectRef>,
 }
 
 unsafe impl Traverse for PyFunction {
@@ -49,6 +52,7 @@ unsafe impl Traverse for PyFunction {
 }
 
 impl PyFunction {
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
         code: PyRef<PyCode>,
         globals: PyDictRef,
@@ -57,6 +61,9 @@ impl PyFunction {
         kw_only_defaults: Option<PyDictRef>,
         qualname: PyStrRef,
         type_params: PyTupleRef,
+        annotations: PyDictRef,
+        module: PyObjectRef,
+        doc: PyObjectRef,
     ) -> Self {
         let name = PyMutex::new(code.obj_name.to_owned());
         PyFunction {
@@ -69,6 +76,9 @@ impl PyFunction {
             type_params: PyMutex::new(type_params),
             #[cfg(feature = "jit")]
             jitted_code: OnceCell::new(),
+            annotations: PyMutex::new(annotations),
+            module: PyMutex::new(module),
+            doc: PyMutex::new(doc),
         }
     }
 
@@ -406,6 +416,41 @@ impl PyFunction {
         *self.name.lock() = name;
     }
 
+    #[pymember(magic)]
+    fn doc(_vm: &VirtualMachine, zelf: PyObjectRef) -> PyResult {
+        let zelf: PyRef<PyFunction> = zelf.downcast().unwrap_or_else(|_| unreachable!());
+        let doc = zelf.doc.lock();
+        Ok(doc.clone())
+    }
+
+    #[pymember(magic, setter)]
+    fn set_doc(vm: &VirtualMachine, zelf: PyObjectRef, value: PySetterValue) -> PyResult<()> {
+        let zelf: PyRef<PyFunction> = zelf.downcast().unwrap_or_else(|_| unreachable!());
+        let value = value.unwrap_or_none(vm);
+        *zelf.doc.lock() = value;
+        Ok(())
+    }
+
+    #[pygetset(magic)]
+    fn module(&self) -> PyObjectRef {
+        self.module.lock().clone()
+    }
+
+    #[pygetset(magic, setter)]
+    fn set_module(&self, module: PySetterValue<PyObjectRef>, vm: &VirtualMachine) {
+        *self.module.lock() = module.unwrap_or_none(vm);
+    }
+
+    #[pygetset(magic)]
+    fn annotations(&self) -> PyDictRef {
+        self.annotations.lock().clone()
+    }
+
+    #[pygetset(magic, setter)]
+    fn set_annotations(&self, annotations: PyDictRef) {
+        *self.annotations.lock() = annotations
+    }
+
     #[pygetset(magic)]
     fn qualname(&self) -> PyStrRef {
         self.qualname.lock().clone()

vm/src/frame.rs

@@ -1731,6 +1731,8 @@ impl ExecutingFrame<'_> {
             None
         };
 
+        let module = vm.unwrap_or_none(self.globals.get_item_opt(identifier!(vm, __name__), vm)?);
+
         // pop argc arguments
         // argument: name, args, globals
         // let scope = self.scope.clone();
@@ -1742,17 +1744,16 @@ impl ExecutingFrame<'_> {
             kw_only_defaults,
             qualified_name.clone(),
             type_params,
+            annotations.downcast().unwrap(),
+            module,
+            vm.ctx.none(),
         )
         .into_pyobject(vm);
 
-        func_obj.set_attr(identifier!(vm, __doc__), vm.ctx.none(), vm)?;
-
         let name = qualified_name.as_str().split('.').next_back().unwrap();
         func_obj.set_attr(identifier!(vm, __name__), vm.new_pyobj(name), vm)?;
         func_obj.set_attr(identifier!(vm, __qualname__), qualified_name, vm)?;
-        let module = vm.unwrap_or_none(self.globals.get_item_opt(identifier!(vm, __name__), vm)?);
-        func_obj.set_attr(identifier!(vm, __module__), module, vm)?;
-        func_obj.set_attr(identifier!(vm, __annotations__), annotations, vm)?;
+        func_obj.set_attr(identifier!(vm, __doc__), vm.ctx.none(), vm)?;
 
         self.push_value(func_obj);
         Ok(None)