add mapping protocol methods to update_slot

Snowapril · Snowapril · commit e0ac7a452ef9 · 2021-09-29T15:24:42.000+09:00
Signed-off-by: snowapril &lt;sinjihng@gmail.com&gt;
diff --git a/vm/src/builtins/pytype.rs b/vm/src/builtins/pytype.rs
@@ -4,13 +4,15 @@ use super::{
 };
 use crate::common::{ascii, lock::PyRwLock};
 use crate::{
+    builtins::dict::PyMapping,
     function::{FuncArgs, KwArgs, OptionalArg},
     slots::{self, Callable, PyTpFlags, PyTypeSlots, SlotGetattro, SlotSetattro},
     utils::Either,
     IdProtocol, PyAttributes, PyClassImpl, PyContext, PyLease, PyObjectRef, PyRef, PyResult,
     PyValue, TryFromObject, TypeProtocol, VirtualMachine,
 };
 use itertools::Itertools;
+use num_traits::ToPrimitive;
 use std::collections::HashSet;
 use std::fmt;
 use std::ops::Deref;
@@ -265,6 +267,54 @@ impl PyType {
                     |zelf, vm| vm.call_special_method(zelf.clone(), "__next__", ());
                 update_slot!(iternext, func);
             }
+            "__len__" | "__getitem__" | "__setitem__" | "__delitem__" => {
+                macro_rules! then_some_closure {
+                    ($cond:expr, $closure:expr) => {
+                        if $cond {
+                            Some($closure)
+                        } else {
+                            None
+                        }
+                    };
+                }
+
+                let func: slots::MappingFunc = |zelf, _vm| {
+                    Ok(PyMapping {
+                        length: then_some_closure!(zelf.has_class_attr("__len__"), |zelf, vm| {
+                            vm.call_special_method(zelf, "__len__", ()).map(|obj| {
+                                obj.payload_if_subclass::<PyInt>(vm)
+                                    .map(|length_obj| {
+                                        length_obj.as_bigint().to_usize().ok_or_else(|| {
+                                            vm.new_value_error(
+                                                "__len__() should return >= 0".to_owned(),
+                                            )
+                                        })
+                                    })
+                                    .unwrap()
+                            })?
+                        }),
+                        subscript: then_some_closure!(
+                            zelf.has_class_attr("__getitem__"),
+                            |zelf: PyObjectRef, needle: PyObjectRef, vm: &VirtualMachine| {
+                                vm.call_special_method(zelf, "__getitem__", (needle,))
+                            }
+                        ),
+                        ass_subscript: then_some_closure!(
+                            zelf.has_class_attr("__setitem__") | zelf.has_class_attr("__delitem__"),
+                            |zelf, needle, value, vm| match value {
+                                Some(value) => vm
+                                    .call_special_method(zelf, "__setitem__", (needle, value),)
+                                    .map(|_| Ok(()))?,
+                                None => vm
+                                    .call_special_method(zelf, "__delitem__", (needle,))
+                                    .map(|_| Ok(()))?,
+                            }
+                        ),
+                    })
+                };
+                update_slot!(as_mapping, func);
+                // TODO: need to update sequence protocol too
+            }
             _ => {}
         }
     }
