From d9c591be9cc3ecbc2b1717f5f0fc82e41dfb58fc Mon Sep 17 00:00:00 2001
From: Bas Schoenmaeckers <b.schoenmaeckers@zandersgroup.com>
Date: Mon, 1 Jun 2026 17:04:31 +0200
Subject: [PATCH] Add number protocol functions

---
 crates/capi/src/abstract_.rs        |  2 +
 crates/capi/src/abstract_/number.rs | 86 +++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+)
 create mode 100644 crates/capi/src/abstract_/number.rs

diff --git a/crates/capi/src/abstract_.rs b/crates/capi/src/abstract_.rs
index 4ef8bbb241a..a677c98c7b4 100644
--- a/crates/capi/src/abstract_.rs
+++ b/crates/capi/src/abstract_.rs
@@ -2,12 +2,14 @@ use crate::{PyObject, pystate::with_vm};
 use alloc::slice;
 use core::ffi::c_int;
 pub use mapping::*;
+pub use number::*;
 use rustpython_vm::builtins::{PyDict, PyStr, PyTuple};
 use rustpython_vm::function::{FuncArgs, KwArgs, PosArgs};
 use rustpython_vm::{AsObject, Py, PyObjectRef, PyResult, VirtualMachine};
 pub use sequence::*;
 
 mod mapping;
+mod number;
 mod sequence;
 
 const PY_VECTORCALL_ARGUMENTS_OFFSET: usize = 1usize << (usize::BITS as usize - 1);
diff --git a/crates/capi/src/abstract_/number.rs b/crates/capi/src/abstract_/number.rs
new file mode 100644
index 00000000000..c5ec492a73d
--- /dev/null
+++ b/crates/capi/src/abstract_/number.rs
@@ -0,0 +1,86 @@
+use crate::{PyObject, pystate::with_vm};
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyNumber_Add(o1: *mut PyObject, o2: *mut PyObject) -> *mut PyObject {
+    with_vm(|vm| vm._add(unsafe { &*o1 }, unsafe { &*o2 }))
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyNumber_Index(obj: *mut PyObject) -> *mut PyObject {
+    with_vm(|vm| unsafe { &*obj }.try_index(vm))
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyNumber_Lshift(o1: *mut PyObject, o2: *mut PyObject) -> *mut PyObject {
+    with_vm(|vm| vm._lshift(unsafe { &*o1 }, unsafe { &*o2 }))
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyNumber_Or(o1: *mut PyObject, o2: *mut PyObject) -> *mut PyObject {
+    with_vm(|vm| vm._or(unsafe { &*o1 }, unsafe { &*o2 }))
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyNumber_Rshift(o1: *mut PyObject, o2: *mut PyObject) -> *mut PyObject {
+    with_vm(|vm| vm._rshift(unsafe { &*o1 }, unsafe { &*o2 }))
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn PyNumber_Subtract(o1: *mut PyObject, o2: *mut PyObject) -> *mut PyObject {
+    with_vm(|vm| vm._sub(unsafe { &*o1 }, unsafe { &*o2 }))
+}
+
+#[cfg(false)]
+mod tests {
+    use pyo3::prelude::*;
+
+    #[test]
+    fn add() {
+        Python::attach(|py| {
+            let lhs = 40i64.into_pyobject(py).unwrap();
+            let rhs = 2i64.into_pyobject(py).unwrap();
+            let out = lhs.add(rhs).unwrap();
+            assert_eq!(out.extract::<i64>().unwrap(), 42);
+        })
+    }
+
+    #[test]
+    fn lshift() {
+        Python::attach(|py| {
+            let lhs = 1i64.into_pyobject(py).unwrap();
+            let rhs = 5i64.into_pyobject(py).unwrap();
+            let out = lhs.lshift(rhs).unwrap();
+            assert_eq!(out.extract::<i64>().unwrap(), 32);
+        })
+    }
+
+    #[test]
+    fn bit_or() {
+        Python::attach(|py| {
+            let lhs = 0b1010i64.into_pyobject(py).unwrap();
+            let rhs = 0b0110i64.into_pyobject(py).unwrap();
+            let out = lhs.bitor(rhs).unwrap();
+            assert_eq!(out.extract::<i64>().unwrap(), 0b1110);
+        })
+    }
+
+    #[test]
+    fn rshift() {
+        Python::attach(|py| {
+            let lhs = 128i64.into_pyobject(py).unwrap();
+            let rhs = 3i64.into_pyobject(py).unwrap();
+            let out = lhs.rshift(rhs).unwrap();
+            assert_eq!(out.extract::<i64>().unwrap(), 16);
+        })
+    }
+
+    #[test]
+    fn subtract() {
+        Python::attach(|py| {
+            let lhs = 50i64.into_pyobject(py).unwrap();
+            let rhs = 8i64.into_pyobject(py).unwrap();
+            let out = lhs.sub(rhs).unwrap();
+            assert_eq!(out.extract::<i64>().unwrap(), 42);
+        })
+    }
+}