From b4124d0d922a8d8c86693f0ae680913d4efa6d60 Mon Sep 17 00:00:00 2001
From: Padraic Fanning <fanninpm@miamioh.edu>
Date: Fri, 3 Mar 2023 18:53:57 -0500
Subject: [PATCH] Add number protocol for PyStr

---
 vm/src/builtins/str.rs | 32 +++++++++++++++++++++++++++-----
 1 file changed, 27 insertions(+), 5 deletions(-)

diff --git a/vm/src/builtins/str.rs b/vm/src/builtins/str.rs
index b9f125894e..1e8778b2c8 100644
--- a/vm/src/builtins/str.rs
+++ b/vm/src/builtins/str.rs
@@ -11,16 +11,16 @@ use crate::{
         format::{FormatSpec, FormatString, FromTemplate},
         str::{BorrowedStr, PyStrKind, PyStrKindData},
     },
-    convert::{IntoPyException, ToPyException, ToPyObject},
+    convert::{IntoPyException, ToPyException, ToPyObject, ToPyResult},
     format::{format, format_map},
     function::{ArgIterable, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue},
     intern::PyInterned,
-    protocol::{PyIterReturn, PyMappingMethods, PySequenceMethods},
+    protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
     sequence::SequenceExt,
     sliceable::{SequenceIndex, SliceableSequenceOp},
     types::{
-        AsMapping, AsSequence, Comparable, Constructor, Hashable, IterNext, IterNextIterable,
-        Iterable, PyComparisonOp, Unconstructible,
+        AsMapping, AsNumber, AsSequence, Comparable, Constructor, Hashable, IterNext,
+        IterNextIterable, Iterable, PyComparisonOp, Unconstructible,
     },
     AsObject, Context, Py, PyExact, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
     TryFromBorrowedObject, VirtualMachine,
@@ -354,7 +354,15 @@ impl PyStr {
 
 #[pyclass(
     flags(BASETYPE),
-    with(AsMapping, AsSequence, Hashable, Comparable, Iterable, Constructor)
+    with(
+        AsMapping,
+        AsNumber,
+        AsSequence,
+        Hashable,
+        Comparable,
+        Iterable,
+        Constructor
+    )
 )]
 impl PyStr {
     #[pymethod(magic)]
@@ -1281,6 +1289,20 @@ impl AsMapping for PyStr {
     }
 }
 
+impl AsNumber for PyStr {
+    fn as_number() -> &'static PyNumberMethods {
+        static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
+            remainder: atomic_func!(|number, other, vm| {
+                PyStr::number_downcast(number)
+                    .modulo(other.to_owned(), vm)
+                    .to_pyresult(vm)
+            }),
+            ..PyNumberMethods::NOT_IMPLEMENTED
+        });
+        &AS_NUMBER
+    }
+}
+
 impl AsSequence for PyStr {
     fn as_sequence() -> &'static PySequenceMethods {
         static AS_SEQUENCE: Lazy<PySequenceMethods> = Lazy::new(|| PySequenceMethods {