Temp: Implement Number Protocol for PyDictKeys and PyDictItems

Author: xiaozhiyan

vm/src/builtins/dict.rs

@@ -1092,7 +1092,15 @@ trait ViewSetOps: DictView {
 }
 
 impl ViewSetOps for PyDictKeys {}
-#[pyclass(with(DictView, Constructor, Comparable, Iterable, ViewSetOps, AsSequence))]
+#[pyclass(with(
+    DictView,
+    Constructor,
+    Comparable,
+    Iterable,
+    ViewSetOps,
+    AsSequence,
+    AsNumber
+))]
 impl PyDictKeys {
     #[pymethod(magic)]
     fn contains(zelf: PyRef<Self>, key: PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
@@ -1133,8 +1141,70 @@ impl AsSequence for PyDictKeys {
     }
 }
 
+impl AsNumber for PyDictKeys {
+    fn as_number() -> &'static PyNumberMethods {
+        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
+            subtract: Some(|num, args, vm| {
+                let num = PySetInner::from_iter(
+                    ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?,
+                    vm,
+                )?;
+                Ok(PySet {
+                    inner: num
+                        .difference(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?,
+                }
+                .into_pyobject(vm))
+            }),
+            and: Some(|num, args, vm| {
+                let num = PySetInner::from_iter(
+                    ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?,
+                    vm,
+                )?;
+                Ok(PySet {
+                    inner: num
+                        .intersection(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?,
+                }
+                .into_pyobject(vm))
+            }),
+            xor: Some(|num, args, vm| {
+                let num = PySetInner::from_iter(
+                    ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?,
+                    vm,
+                )?;
+                Ok(PySet {
+                    inner: num.symmetric_difference(
+                        ArgIterable::try_from_object(vm, args.to_owned())?,
+                        vm,
+                    )?,
+                }
+                .into_pyobject(vm))
+            }),
+            or: Some(|num, args, vm| {
+                let num = PySetInner::from_iter(
+                    ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?,
+                    vm,
+                )?;
+                Ok(PySet {
+                    inner: num.union(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?,
+                }
+                .into_pyobject(vm))
+            }),
+            ..PyNumberMethods::NOT_IMPLEMENTED
+        };
+        &AS_NUMBER
+    }
+}
+
 impl ViewSetOps for PyDictItems {}
-#[pyclass(with(DictView, Constructor, Comparable, Iterable, ViewSetOps, AsSequence))]
+#[pyclass(with(
+    DictView,
+    Constructor,
+    Comparable,
+    Iterable,
+    ViewSetOps,
+    AsSequence,
+    AsNumber
+))]
 impl PyDictItems {
     #[pymethod(magic)]
     fn contains(zelf: PyRef<Self>, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
@@ -1189,6 +1259,60 @@ impl AsSequence for PyDictItems {
     }
 }
 
+impl AsNumber for PyDictItems {
+    fn as_number() -> &'static PyNumberMethods {
+        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
+            subtract: Some(|num, args, vm| {
+                let num = PySetInner::from_iter(
+                    ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?,
+                    vm,
+                )?;
+                Ok(PySet {
+                    inner: num
+                        .difference(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?,
+                }
+                .into_pyobject(vm))
+            }),
+            and: Some(|num, args, vm| {
+                let num = PySetInner::from_iter(
+                    ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?,
+                    vm,
+                )?;
+                Ok(PySet {
+                    inner: num
+                        .intersection(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?,
+                }
+                .into_pyobject(vm))
+            }),
+            xor: Some(|num, args, vm| {
+                let num = PySetInner::from_iter(
+                    ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?,
+                    vm,
+                )?;
+                Ok(PySet {
+                    inner: num.symmetric_difference(
+                        ArgIterable::try_from_object(vm, args.to_owned())?,
+                        vm,
+                    )?,
+                }
+                .into_pyobject(vm))
+            }),
+            or: Some(|num, args, vm| {
+                let num = PySetInner::from_iter(
+                    ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?,
+                    vm,
+                )?;
+                Ok(PySet {
+                    inner: num.union(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?,
+                }
+                .into_pyobject(vm))
+            }),
+            ..PyNumberMethods::NOT_IMPLEMENTED
+        };
+        &AS_NUMBER
+    }
+}
+
 #[pyclass(with(DictView, Constructor, Iterable, AsSequence))]
 impl PyDictValues {
     #[pygetset]