diff --git a/vm/src/builtins/function.rs b/vm/src/builtins/function.rs index ab069e0e7..e4cc49f9d 100644 --- a/vm/src/builtins/function.rs +++ b/vm/src/builtins/function.rs @@ -38,6 +38,9 @@ pub struct PyFunction { type_params: PyMutex, #[cfg(feature = "jit")] jitted_code: OnceCell, + annotations: PyMutex, + module: PyMutex, + doc: PyMutex, } 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, globals: PyDictRef, @@ -57,6 +61,9 @@ impl PyFunction { kw_only_defaults: Option, 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 = 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 = 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, 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() diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 606ab8300..5a42a1f88 100644 --- a/vm/src/frame.rs +++ b/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)