another fix

youknowone · youknowone · commit 3ab63ef5398f · 2025-07-21T09:36:47.000+09:00
diff --git a/vm/src/builtins/type.rs b/vm/src/builtins/type.rs
@@ -1446,63 +1446,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(())
     }
 }
 
