From e8755a75c654cc2cc95dcf866e886ac75c1f3c92 Mon Sep 17 00:00:00 2001
From: "Kim, YeonWoo" <oow214@gmail.com>
Date: Sat, 22 Oct 2022 01:50:58 +0900
Subject: [PATCH] Add set annotations for type

---
 vm/src/builtins/type.rs | 28 ++++++++++++++++++++++++++++
 vm/src/types/slot.rs    |  1 +
 2 files changed, 29 insertions(+)

diff --git a/vm/src/builtins/type.rs b/vm/src/builtins/type.rs
index 3758d3dfd..9a2f21b31 100644
--- a/vm/src/builtins/type.rs
+++ b/vm/src/builtins/type.rs
@@ -549,6 +549,34 @@ impl PyType {
         Ok(annotations)
     }
 
+    #[pygetset(magic, setter)]
+    fn set_annotations(&self, value: Option<PyObjectRef>, vm: &VirtualMachine) -> PyResult<()> {
+        if self.slots.flags.has_feature(PyTypeFlags::IMMUTABLETYPE) {
+            return Err(vm.new_type_error(format!(
+                "cannot set '__annotations__' attribute of immutable type '{}'",
+                self.name()
+            )));
+        }
+
+        let __annotations__ = identifier!(vm, __annotations__);
+        if let Some(value) = value {
+            self.attributes.write().insert(__annotations__, value);
+        } else {
+            self.attributes
+                .read()
+                .get(__annotations__)
+                .cloned()
+                .ok_or_else(|| {
+                    vm.new_attribute_error(format!(
+                        "'{}' object has no attribute '__annotations__'",
+                        self.name()
+                    ))
+                })?;
+        }
+
+        Ok(())
+    }
+
     #[pygetset(magic)]
     pub fn module(&self, vm: &VirtualMachine) -> PyObjectRef {
         self.attributes
diff --git a/vm/src/types/slot.rs b/vm/src/types/slot.rs
index 57399388b..e1578c5f7 100644
--- a/vm/src/types/slot.rs
+++ b/vm/src/types/slot.rs
@@ -105,6 +105,7 @@ impl std::fmt::Debug for PyTypeSlots {
 bitflags! {
     #[non_exhaustive]
     pub struct PyTypeFlags: u64 {
+        const IMMUTABLETYPE = 1 << 8;
         const HEAPTYPE = 1 << 9;
         const BASETYPE = 1 << 10;
         const METHOD_DESCR = 1 << 17;