another fix

youknowone · youknowone · commit 6930d0c30c92 · 2025-07-21T14:28:21.000+09:00
diff --git a/vm/src/builtins/type.rs b/vm/src/builtins/type.rs
@@ -1019,7 +1019,7 @@ impl Constructor for PyType {
             attributes.insert(identifier!(vm, __hash__), vm.ctx.none.clone().into());
         }
 
-        let heaptype_slots: Option<PyRef<PyTuple<PyStrRef>>> =
+        let (heaptype_slots, add_dict): (Option<PyRef<PyTuple<PyStrRef>>>, bool) =
             if let Some(x) = attributes.get(identifier!(vm, __slots__)) {
                 let slots = if x.class().is(vm.ctx.types.str_type) {
                     let x = unsafe { x.downcast_unchecked_ref::<PyStr>() };
@@ -1036,9 +1036,26 @@ impl Constructor for PyType {
                     let tuple = elements.into_pytuple(vm);
                     tuple.try_into_typed(vm)?
                 };
-                Some(slots)
+
+                // Check if __dict__ is in slots
+                let dict_name = "__dict__";
+                let has_dict = slots.iter().any(|s| s.as_str() == dict_name);
+
+                // Filter out __dict__ from slots
+                let filtered_slots = if has_dict {
+                    let filtered: Vec<PyStrRef> = slots
+                        .iter()
+                        .filter(|s| s.as_str() != dict_name)
+                        .cloned()
+                        .collect();
+                    PyTuple::new_ref_typed(filtered, &vm.ctx)
+                } else {
+                    slots
+                };
+
+                (Some(filtered_slots), has_dict)
             } else {
-                None
+                (None, false)
             };
 
         // FIXME: this is a temporary fix. multi bases with multiple slots will break object
@@ -1051,8 +1068,10 @@ impl Constructor for PyType {
         let member_count: usize = base_member_count + heaptype_member_count;
 
         let mut flags = PyTypeFlags::heap_type_flags();
-        // Only add HAS_DICT and MANAGED_DICT if __slots__ is not defined.
-        if heaptype_slots.is_none() {
+        // Add HAS_DICT and MANAGED_DICT if:
+        // 1. __slots__ is not defined, OR
+        // 2. __dict__ is in __slots__
+        if heaptype_slots.is_none() || add_dict {
             flags |= PyTypeFlags::HAS_DICT | PyTypeFlags::MANAGED_DICT;
         }
 
@@ -1446,63 +1465,85 @@ impl Representable for PyType {
     }
 }
 
-fn find_base_dict_descr(cls: &Py<PyType>, vm: &VirtualMachine) -> Option<PyObjectRef> {
-    let dict_attr = identifier!(vm, __dict__);
-
-    // First check if the class itself has a user-defined __dict__ property
-    if let Some(descr) = cls.attributes.read().get(&dict_attr) {
-        // If it's a property (user-defined), return it
-        if descr.class().is(vm.ctx.types.property_type) {
-            return Some(descr.clone());
+// Equivalent to CPython's get_builtin_base_with_dict
+fn get_builtin_base_with_dict(typ: &Py<PyType>, vm: &VirtualMachine) -> Option<PyTypeRef> {
+    let mut current = Some(typ.to_owned());
+    while let Some(t) = current {
+        // In CPython: type->tp_dictoffset != 0 && !(type->tp_flags & Py_TPFLAGS_HEAPTYPE)
+        // Special case: type itself is a builtin with dict support
+        if t.is(vm.ctx.types.type_type) {
+            return Some(t);
+        }
+        // We check HAS_DICT flag (equivalent to tp_dictoffset != 0) and HEAPTYPE
+        if t.slots.flags.contains(PyTypeFlags::HAS_DICT)
+            && !t.slots.flags.contains(PyTypeFlags::HEAPTYPE)
+        {
+            return Some(t);
         }
-        // Skip getset descriptors in the current class to avoid recursion
+        current = t.__base__();
     }
+    None
+}
 
-    // Then check bases (like original implementation)
-    cls.iter_base_chain().skip(1).find_map(|cls| {
-        // TODO: should actually be some translation of:
-        // cls.slot_dictoffset != 0 && !cls.flags.contains(HEAPTYPE)
-        if cls.is(vm.ctx.types.type_type) {
-            cls.get_attr(dict_attr)
-        } else {
-            None
-        }
-    })
+// Equivalent to CPython's get_dict_descriptor
+fn get_dict_descriptor(base: &Py<PyType>, vm: &VirtualMachine) -> Option<PyObjectRef> {
+    let dict_attr = identifier!(vm, __dict__);
+    // Use _PyType_Lookup (which is lookup_ref in RustPython)
+    base.lookup_ref(dict_attr, vm)
 }
 
 fn subtype_get_dict(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-    // TODO: obj.class().as_pyref() need to be supported
-    let ret = match find_base_dict_descr(obj.class(), vm) {
-        Some(descr) => vm.call_get_descriptor(&descr, obj).unwrap_or_else(|| {
+    let base = get_builtin_base_with_dict(obj.class(), vm);
+
+    if let Some(base_type) = base {
+        if let Some(descr) = get_dict_descriptor(&base_type, vm) {
+            // Call the descriptor's tp_descr_get
+            vm.call_get_descriptor(&descr, obj.clone())
+                .unwrap_or_else(|| {
+                    Err(vm.new_type_error(format!(
+                        "this __dict__ descriptor does not support '{}' objects",
+                        obj.class().name()
+                    )))
+                })
+        } else {
             Err(vm.new_type_error(format!(
                 "this __dict__ descriptor does not support '{}' objects",
-                descr.class()
+                obj.class().name()
             )))
-        })?,
-        None => object::object_get_dict(obj, vm)?.into(),
-    };
-    Ok(ret)
+        }
+    } else {
+        // PyObject_GenericGetDict
+        object::object_get_dict(obj, vm).map(Into::into)
+    }
 }
 
 fn subtype_set_dict(obj: PyObjectRef, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
-    let cls = obj.class();
-    match find_base_dict_descr(cls, vm) {
-        Some(descr) => {
+    // Following CPython's subtype_setdict exactly
+    let base = get_builtin_base_with_dict(obj.class(), vm);
+
+    if let Some(base_type) = base {
+        if let Some(descr) = get_dict_descriptor(&base_type, vm) {
+            // Call the descriptor's tp_descr_set
             let descr_set = descr
                 .class()
                 .mro_find_map(|cls| cls.slots.descr_set.load())
                 .ok_or_else(|| {
                     vm.new_type_error(format!(
                         "this __dict__ descriptor does not support '{}' objects",
-                        cls.name()
+                        obj.class().name()
                     ))
                 })?;
             descr_set(&descr, obj, PySetterValue::Assign(value), vm)
+        } else {
+            Err(vm.new_type_error(format!(
+                "this __dict__ descriptor does not support '{}' objects",
+                obj.class().name()
+            )))
         }
-        None => {
-            object::object_set_dict(obj, value.try_into_value(vm)?, vm)?;
-            Ok(())
-        }
+    } else {
+        // PyObject_GenericSetDict
+        object::object_set_dict(obj, value.try_into_value(vm)?, vm)?;
+        Ok(())
     }
 }
 
