From 5f0a58fbc6f6db53eb6ea8745dbefc044341bb4a Mon Sep 17 00:00:00 2001
From: Kangzhi Shi <shikangzhi@gmail.com>
Date: Wed, 17 Nov 2021 09:38:48 +0200
Subject: [PATCH] impl sequence protocol for mappingproxy

---
 vm/src/builtins/mappingproxy.rs | 42 +++++++++++++++++++++++++--------
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/vm/src/builtins/mappingproxy.rs b/vm/src/builtins/mappingproxy.rs
index a650fffd2..05b232b0c 100644
--- a/vm/src/builtins/mappingproxy.rs
+++ b/vm/src/builtins/mappingproxy.rs
@@ -1,4 +1,6 @@
-use super::{PyDict, PyGenericAlias, PyList, PyStrRef, PyTuple, PyTypeRef};
+use std::borrow::Cow;
+
+use super::{PyDict, PyGenericAlias, PyList, PyStr, PyStrRef, PyTuple, PyTypeRef};
 use crate::{
     function::{IntoPyObject, OptionalArg},
     protocol::{PyMapping, PyMappingMethods},
@@ -54,7 +56,7 @@ impl Constructor for PyMappingProxy {
     }
 }
 
-#[pyimpl(with(AsMapping, Iterable, Constructor))]
+#[pyimpl(with(AsMapping, Iterable, Constructor, AsSequence))]
 impl PyMappingProxy {
     fn get_inner(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult<Option<PyObjectRef>> {
         let opt = match &self.mapping {
@@ -85,20 +87,24 @@ impl PyMappingProxy {
             .ok_or_else(|| vm.new_key_error(key))
     }
 
-    #[pymethod(magic)]
-    pub fn contains(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+    fn _contains(&self, key: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
         match &self.mapping {
             MappingProxyInner::Class(class) => {
-                let key = PyStrRef::try_from_object(vm, key)?;
-                Ok(vm
-                    .ctx
-                    .new_bool(class.attributes.read().contains_key(key.as_str()))
-                    .into())
+                // let key = PyStrRef::try_from_object(vm, key)?;
+                let key = key
+                    .payload::<PyStr>()
+                    .ok_or_else(|| pyref_type_error(vm, PyStr::class(vm), key))?;
+                Ok(class.attributes.read().contains_key(key.as_str()))
             }
-            MappingProxyInner::Dict(obj) => vm._membership(obj.clone(), key),
+            MappingProxyInner::Dict(obj) => PySequence::from(obj.as_ref()).contains(key, vm),
         }
     }
 
+    #[pymethod(magic)]
+    pub fn contains(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
+        self._contains(&key, vm)
+    }
+
     #[pymethod]
     pub fn items(&self, vm: &VirtualMachine) -> PyResult {
         let obj = match &self.mapping {
@@ -171,6 +177,22 @@ impl AsMapping for PyMappingProxy {
     }
 }
 
+impl AsSequence for PyMappingProxy {
+    fn as_sequence(
+        _zelf: &crate::PyObjectView<Self>,
+        _vm: &VirtualMachine,
+    ) -> Cow<'static, PySequenceMethods> {
+        Cow::Borrowed(&Self::SEQUENCE_METHODS)
+    }
+}
+
+impl PyMappingProxy {
+    const SEQUENCE_METHODS: PySequenceMethods = PySequenceMethods {
+        contains: Some(|seq, target, vm| seq.obj_as::<Self>()._contains(target, vm)),
+        ..*PySequenceMethods::not_implemented()
+    };
+}
+
 impl Iterable for PyMappingProxy {
     fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
         let obj = match &zelf.mapping {