Merge pull request RustPython#4619 from xiaozhiyan/implement-number-protocol-for-pytype-pyset-pyfrozenset

Implement Number Protocol for `PyType`, `PySet` and `PyFrozenSet`

Author: youknowone

vm/src/builtins/set.rs

@@ -1,20 +1,20 @@
-use once_cell::sync::Lazy;
-
 /*
  * Builtin set type with a sequence of unique items.
  */
 use super::{
     builtins_iter, IterStatus, PositionIterInternal, PyDict, PyDictRef, PyGenericAlias, PyTupleRef,
     PyType, PyTypeRef,
 };
-use crate::atomic_func;
-use crate::common::{ascii, hash::PyHash, lock::PyMutex, rc::PyRc};
 use crate::{
+    atomic_func,
     class::PyClassImpl,
+    common::{ascii, hash::PyHash, lock::PyMutex, rc::PyRc},
+    convert::ToPyResult,
     dictdatatype::{self, DictSize},
     function::{ArgIterable, FuncArgs, OptionalArg, PosArgs, PyArithmeticValue, PyComparisonValue},
-    protocol::{PyIterReturn, PySequenceMethods},
+    protocol::{PyIterReturn, PyNumberMethods, PySequenceMethods},
     recursion::ReprGuard,
+    types::AsNumber,
     types::{
         AsSequence, Comparable, Constructor, Hashable, Initializer, IterNext, IterNextIterable,
         Iterable, PyComparisonOp, Unconstructible, Unhashable,
@@ -23,6 +23,7 @@ use crate::{
     vm::VirtualMachine,
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
 };
+use once_cell::sync::Lazy;
 use std::{fmt, ops::Deref};
 
 pub type SetContentType = dictdatatype::Dict<()>;
@@ -488,7 +489,15 @@ fn reduce_set(
 }
 
 #[pyclass(
-    with(Constructor, Initializer, AsSequence, Hashable, Comparable, Iterable),
+    with(
+        Constructor,
+        Initializer,
+        AsSequence,
+        Hashable,
+        Comparable,
+        Iterable,
+        AsNumber
+    ),
     flags(BASETYPE)
 )]
 impl PySet {
@@ -792,6 +801,67 @@ impl Iterable for PySet {
     }
 }
 
+impl AsNumber for PySet {
+    fn as_number() -> &'static PyNumberMethods {
+        static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
+            subtract: atomic_func!(|number, other, vm| {
+                PySet::number_downcast(number)
+                    .sub(other.to_owned(), vm)
+                    .to_pyresult(vm)
+            }),
+            and: atomic_func!(|number, other, vm| {
+                PySet::number_downcast(number)
+                    .and(other.to_owned(), vm)
+                    .to_pyresult(vm)
+            }),
+            xor: atomic_func!(|number, other, vm| {
+                PySet::number_downcast(number)
+                    .xor(other.to_owned(), vm)
+                    .to_pyresult(vm)
+            }),
+            or: atomic_func!(|number, other, vm| {
+                PySet::number_downcast(number)
+                    .or(other.to_owned(), vm)
+                    .to_pyresult(vm)
+            }),
+            inplace_subtract: atomic_func!(|number, other, vm| {
+                PySet::isub(
+                    PySet::number_downcast(number).to_owned(),
+                    AnySet::try_from_object(vm, other.to_owned())?,
+                    vm,
+                )
+                .to_pyresult(vm)
+            }),
+            inplace_and: atomic_func!(|number, other, vm| {
+                PySet::iand(
+                    PySet::number_downcast(number).to_owned(),
+                    AnySet::try_from_object(vm, other.to_owned())?,
+                    vm,
+                )
+                .to_pyresult(vm)
+            }),
+            inplace_xor: atomic_func!(|number, other, vm| {
+                PySet::ixor(
+                    PySet::number_downcast(number).to_owned(),
+                    AnySet::try_from_object(vm, other.to_owned())?,
+                    vm,
+                )
+                .to_pyresult(vm)
+            }),
+            inplace_or: atomic_func!(|number, other, vm| {
+                PySet::ior(
+                    PySet::number_downcast(number).to_owned(),
+                    AnySet::try_from_object(vm, other.to_owned())?,
+                    vm,
+                )
+                .to_pyresult(vm)
+            }),
+            ..PyNumberMethods::NOT_IMPLEMENTED
+        });
+        &AS_NUMBER
+    }
+}
+
 impl Constructor for PyFrozenSet {
     type Args = OptionalArg<PyObjectRef>;
 
@@ -822,7 +892,7 @@ impl Constructor for PyFrozenSet {
 
 #[pyclass(
     flags(BASETYPE),
-    with(Constructor, AsSequence, Hashable, Comparable, Iterable)
+    with(Constructor, AsSequence, Hashable, Comparable, Iterable, AsNumber)
 )]
 impl PyFrozenSet {
     #[pymethod(magic)]
@@ -1019,6 +1089,35 @@ impl Iterable for PyFrozenSet {
     }
 }
 
+impl AsNumber for PyFrozenSet {
+    fn as_number() -> &'static PyNumberMethods {
+        static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
+            subtract: atomic_func!(|number, other, vm| {
+                PyFrozenSet::number_downcast(number)
+                    .sub(other.to_owned(), vm)
+                    .to_pyresult(vm)
+            }),
+            and: atomic_func!(|number, other, vm| {
+                PyFrozenSet::number_downcast(number)
+                    .and(other.to_owned(), vm)
+                    .to_pyresult(vm)
+            }),
+            xor: atomic_func!(|number, other, vm| {
+                PyFrozenSet::number_downcast(number)
+                    .xor(other.to_owned(), vm)
+                    .to_pyresult(vm)
+            }),
+            or: atomic_func!(|number, other, vm| {
+                PyFrozenSet::number_downcast(number)
+                    .or(other.to_owned(), vm)
+                    .to_pyresult(vm)
+            }),
+            ..PyNumberMethods::NOT_IMPLEMENTED
+        });
+        &AS_NUMBER
+    }
+}
+
 struct AnySet {
     object: PyObjectRef,
 }

vm/src/builtins/type.rs

@@ -2,12 +2,8 @@ use super::{
     mappingproxy::PyMappingProxy, object, union_, PyClassMethod, PyDictRef, PyList, PyStaticMethod,
     PyStr, PyStrInterned, PyStrRef, PyTuple, PyTupleRef, PyWeak,
 };
-use crate::common::{
-    ascii,
-    borrow::BorrowedValue,
-    lock::{PyRwLock, PyRwLockReadGuard},
-};
 use crate::{
+    atomic_func,
     builtins::{
         descriptor::{
             DescrObject, MemberDef, MemberDescrObject, MemberGetter, MemberKind, MemberSetter,
@@ -17,15 +13,22 @@ use crate::{
         PyBaseExceptionRef,
     },
     class::{PyClassImpl, StaticType},
-    convert::ToPyObject,
+    common::{
+        ascii,
+        borrow::BorrowedValue,
+        lock::{PyRwLock, PyRwLockReadGuard},
+    },
+    convert::{ToPyObject, ToPyResult},
     function::{FuncArgs, KwArgs, OptionalArg, PySetterValue},
     identifier,
     protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
+    types::AsNumber,
     types::{Callable, GetAttr, PyTypeFlags, PyTypeSlots, SetAttr},
     AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine,
 };
 use indexmap::{map::Entry, IndexMap};
 use itertools::Itertools;
+use once_cell::sync::Lazy;
 use std::{borrow::Borrow, collections::HashSet, fmt, ops::Deref, pin::Pin, ptr::NonNull};
 
 #[pyclass(module = false, name = "type")]
@@ -371,7 +374,7 @@ impl Py<PyType> {
     }
 }
 
-#[pyclass(with(GetAttr, SetAttr, Callable), flags(BASETYPE))]
+#[pyclass(with(GetAttr, SetAttr, Callable, AsNumber), flags(BASETYPE))]
 impl PyType {
     // bound method for every type
     pub(crate) fn __new__(zelf: PyRef<Self>, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
@@ -1058,6 +1061,18 @@ impl Callable for PyType {
     }
 }
 
+impl AsNumber for PyType {
+    fn as_number() -> &'static PyNumberMethods {
+        static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
+            or: atomic_func!(|num, other, vm| {
+                or_(num.obj.to_owned(), other.to_owned(), vm).to_pyresult(vm)
+            }),
+            ..PyNumberMethods::NOT_IMPLEMENTED
+        });
+        &AS_NUMBER
+    }
+}
+
 fn find_base_dict_descr(cls: &Py<PyType>, vm: &VirtualMachine) -> Option<PyObjectRef> {
     cls.iter_base_chain().skip(1).find_map(|cls| {
         // TODO: should actually be some translation of: