temporary fill up missing number protocols

youknowone · youknowone · commit 70b576ee03ef · 2023-03-09T19:50:13.000+09:00
diff --git a/derive-impl/src/pyclass.rs b/derive-impl/src/pyclass.rs
@@ -548,14 +548,27 @@ where
                     other
                 ),
             };
-            quote_spanned! { ident.span() =>
-                class.set_str_attr(
-                    #py_name,
-                    ctx.make_funcdef(#py_name, Self::#ident)
-                        #doc
-                        #build_func,
-                    ctx,
-                );
+            if py_name.starts_with("__") && py_name.ends_with("__") {
+                let name_ident = Ident::new(&py_name, ident.span());
+                quote_spanned! { ident.span() =>
+                    class.set_attr(
+                        ctx.names.#name_ident,
+                        ctx.make_funcdef(#py_name, Self::#ident)
+                            #doc
+                            #build_func
+                        .into(),
+                    );
+                }
+            } else {
+                quote_spanned! { ident.span() =>
+                    class.set_str_attr(
+                        #py_name,
+                        ctx.make_funcdef(#py_name, Self::#ident)
+                            #doc
+                            #build_func,
+                        ctx,
+                    );
+                }
             }
         };
 
diff --git a/stdlib/src/array.rs b/stdlib/src/array.rs
@@ -6,6 +6,7 @@ pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let array = module
         .get_attr("array", vm)
         .expect("Expect array has array type.");
+    array.init_builtin_number_slots(&vm.ctx);
 
     let collections_abc = vm
         .import("collections.abc", None, 0)
diff --git a/vm/src/builtins/type.rs b/vm/src/builtins/type.rs
@@ -187,25 +187,6 @@ impl PyType {
 
         *slots.name.get_mut() = Some(String::from(name));
 
-        #[allow(clippy::mutable_key_type)]
-        let mut slot_name_set = HashSet::new();
-
-        for cls in mro.iter() {
-            for &name in cls.attributes.read().keys() {
-                if name != identifier!(ctx, __new__)
-                    && name.as_str().starts_with("__")
-                    && name.as_str().ends_with("__")
-                {
-                    slot_name_set.insert(name);
-                }
-            }
-        }
-        for &name in attrs.keys() {
-            if name.as_str().starts_with("__") && name.as_str().ends_with("__") {
-                slot_name_set.insert(name);
-            }
-        }
-
         let new_type = PyRef::new_ref(
             PyType {
                 base: Some(base),
@@ -220,9 +201,7 @@ impl PyType {
             None,
         );
 
-        for attr_name in slot_name_set {
-            new_type.update_slot::<true>(attr_name, ctx);
-        }
+        new_type.init_slots(ctx);
 
         let weakref_type = super::PyWeak::static_type();
         for base in &new_type.bases {
@@ -280,6 +259,30 @@ impl PyType {
         Ok(new_type)
     }
 
+    pub(crate) fn init_slots(&self, ctx: &Context) {
+        #[allow(clippy::mutable_key_type)]
+        let mut slot_name_set = std::collections::HashSet::new();
+
+        for cls in self.mro.iter() {
+            for &name in cls.attributes.read().keys() {
+                if name == identifier!(ctx, __new__) {
+                    continue;
+                }
+                if name.as_str().starts_with("__") && name.as_str().ends_with("__") {
+                    slot_name_set.insert(name);
+                }
+            }
+        }
+        for &name in self.attributes.read().keys() {
+            if name.as_str().starts_with("__") && name.as_str().ends_with("__") {
+                slot_name_set.insert(name);
+            }
+        }
+        for attr_name in slot_name_set {
+            self.update_slot::<true>(attr_name, ctx);
+        }
+    }
+
     pub fn slot_name(&self) -> String {
         self.slots.name.read().as_ref().unwrap().to_string()
     }
@@ -1329,3 +1332,64 @@ mod tests {
         );
     }
 }
+
+impl crate::PyObject {
+    // temporary tool to fill missing number protocols for builtin types
+    pub fn init_builtin_number_slots(&self, ctx: &Context) {
+        let typ = self
+            .downcast_ref::<PyType>()
+            .expect("not called from a type");
+        macro_rules! call_update_slot {
+            ($name:ident) => {
+                let id = identifier!(ctx, $name);
+                if typ.has_attr(id) {
+                    typ.update_slot::<true>(identifier!(ctx, $name), ctx);
+                }
+            };
+        }
+        call_update_slot!(__add__);
+        call_update_slot!(__radd__);
+        call_update_slot!(__iadd__);
+        call_update_slot!(__sub__);
+        call_update_slot!(__rsub__);
+        call_update_slot!(__isub__);
+        call_update_slot!(__mul__);
+        call_update_slot!(__rmul__);
+        call_update_slot!(__imul__);
+        call_update_slot!(__mod__);
+        call_update_slot!(__rmod__);
+        call_update_slot!(__imod__);
+        call_update_slot!(__div__);
+        call_update_slot!(__rdiv__);
+        call_update_slot!(__idiv__);
+        call_update_slot!(__divmod__);
+        call_update_slot!(__rdivmod__);
+        call_update_slot!(__pow__);
+        call_update_slot!(__rpow__);
+        call_update_slot!(__ipow__);
+        call_update_slot!(__lshift__);
+        call_update_slot!(__rlshift__);
+        call_update_slot!(__ilshift__);
+        call_update_slot!(__rshift__);
+        call_update_slot!(__rrshift__);
+        call_update_slot!(__irshift__);
+        call_update_slot!(__and__);
+        call_update_slot!(__rand__);
+        call_update_slot!(__iand__);
+        call_update_slot!(__xor__);
+        call_update_slot!(__rxor__);
+        call_update_slot!(__ixor__);
+        call_update_slot!(__or__);
+        call_update_slot!(__ror__);
+        call_update_slot!(__ior__);
+        call_update_slot!(__floordiv__);
+        call_update_slot!(__rfloordiv__);
+        call_update_slot!(__ifloordiv__);
+        call_update_slot!(__truediv__);
+        call_update_slot!(__rtruediv__);
+        call_update_slot!(__itruediv__);
+        call_update_slot!(__matmul__);
+        call_update_slot!(__rmatmul__);
+        call_update_slot!(__imatmul__);
+    }
+}
diff --git a/vm/src/stdlib/builtins.rs b/vm/src/stdlib/builtins.rs
@@ -957,6 +957,11 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) {
     crate::protocol::VecBuffer::make_class(&vm.ctx);
 
     builtins::extend_module(vm, &module);
+    use crate::AsObject;
+    ctx.types
+        .generic_alias_type
+        .as_object()
+        .init_builtin_number_slots(&vm.ctx);
 
     let debug_mode: bool = vm.state.settings.optimize == 0;
     extend_module!(vm, module, {
diff --git a/vm/src/vm/context.rs b/vm/src/vm/context.rs
@@ -78,6 +78,7 @@ declare_const_name! {
     __aenter__,
     __aexit__,
     __aiter__,
+    __alloc__,
     __all__,
     __and__,
     __anext__,
@@ -121,7 +122,9 @@ declare_const_name! {
     __get__,
     __getattr__,
     __getattribute__,
+    __getformat__,
     __getitem__,
+    __getnewargs__,
     __gt__,
     __hash__,
     __iadd__,
@@ -146,6 +149,7 @@ declare_const_name! {
     __iter__,
     __itruediv__,
     __ixor__,
+    __jit__,  // RustPython dialect
     __le__,
     __len__,
     __length_hint__,
@@ -195,13 +199,16 @@ declare_const_name! {
     __rtruediv__,
     __rxor__,
     __set__,
-    __set_name__,
     __setattr__,
     __setitem__,
+    __setstate__,
+    __set_name__,
     __slots__,
     __str__,
     __sub__,
     __subclasscheck__,
+    __subclasshook__,
+    __subclasses__,
     __sizeof__,
     __truediv__,
     __trunc__,
