Merge pull request RustPython#4574 from xiaozhiyan/implement-number-protocol-for-pybytearray

youknowone · web-flow · commit 204643b8da44 · 2023-02-25T22:58:24.000+09:00
Implement Number Protocol for `PyByteArray`
diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs
@@ -19,24 +19,26 @@ use crate::{
             PyRwLockReadGuard, PyRwLockWriteGuard,
         },
     },
-    convert::ToPyObject,
+    convert::{ToPyObject, ToPyResult},
     function::Either,
     function::{
         ArgBytesLike, ArgIterable, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue,
     },
     protocol::{
         BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
-        PyMappingMethods, PySequenceMethods,
+        PyMappingMethods, PyNumberMethods, PySequenceMethods,
     },
     sliceable::{SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp},
     types::{
-        AsBuffer, AsMapping, AsSequence, Callable, Comparable, Constructor, Hashable, Initializer,
-        IterNext, IterNextIterable, Iterable, PyComparisonOp, Unconstructible, Unhashable,
+        AsBuffer, AsMapping, AsNumber, AsSequence, Callable, Comparable, Constructor, Hashable,
+        Initializer, IterNext, IterNextIterable, Iterable, PyComparisonOp, Unconstructible,
+        Unhashable,
     },
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
     VirtualMachine,
 };
 use bstr::ByteSlice;
+use once_cell::sync::Lazy;
 use std::mem::size_of;
 
 #[pyclass(module = false, name = "bytearray")]
@@ -103,6 +105,7 @@ pub(crate) fn init(context: &Context) {
         AsBuffer,
         AsMapping,
         AsSequence,
+        AsNumber,
         Iterable
     )
 )]
@@ -856,6 +859,20 @@ impl AsSequence for PyByteArray {
     }
 }
 
+impl AsNumber for PyByteArray {
+    fn as_number() -> &'static PyNumberMethods {
+        static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
+            remainder: atomic_func!(|number, other, vm| {
+                PyByteArray::number_downcast(number)
+                    .mod_(other.to_owned(), vm)
+                    .to_pyresult(vm)
+            }),
+            ..PyNumberMethods::NOT_IMPLEMENTED
+        });
+        &AS_NUMBER
+    }
+}
+
 impl Unhashable for PyByteArray {}
 
 impl Iterable for PyByteArray {
