Skip to content

Commit 0f769c4

Browse files
dgrunwalddgrunwald
committed
Merge branch 'feature/number-protocol' of https://github.com/Techcable/rust-cpython 
2 parents a8b2496 + 830fbed commit 0f769c4

File tree

3 files changed

+328
-1
lines changed

3 files changed

+328
-1
lines changed

‎src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ pub use ffi::Py_ssize_t;
101101

102102
pub use crate::conversion::{FromPyObject, RefFromPyObject, ToPyObject};
103103
pub use crate::err::{PyErr, PyResult};
104-
pub use crate::objectprotocol::ObjectProtocol;
104+
pub use crate::objectprotocol::{ObjectProtocol, NumberProtocol};
105105
pub use crate::objects::*;
106106
pub use crate::py_class::CompareOp;
107107
pub use crate::python::{

‎src/objectprotocol.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@ use crate::ffi;
2525
use crate::objects::{PyDict, PyObject, PyString, PyTuple};
2626
use crate::python::{Python, PythonObject, ToPythonPointer};
2727

28+
mod number;
29+
30+
pub use self::number::NumberProtocol;
31+
2832
/// Trait that contains methods
2933
pub trait ObjectProtocol: PythonObject {
3034
/// Determines whether this object has the given attribute.

‎src/objectprotocol/number.rs‎

Lines changed: 323 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,323 @@
1+
use std::cmp::Ordering;
2+
use std::fmt;
3+
4+
use crate::conversion::ToPyObject;
5+
use crate::err::{self, PyErr, PyResult};
6+
use crate::ffi;
7+
use crate::objects::{PyObject, PyInt, PyLong, PyFloat};
8+
use crate::python::{Python, PythonObject, ToPythonPointer};
9+
10+
11+
use super::ObjectProtocol;
12+
13+
/// Operations on numeric objects
14+
pub trait NumberProtocol: ObjectProtocol {
15+
/// Perform addition (self + other)
16+
///
17+
/// Invokes the `__add__` magic-method
18+
#[inline]
19+
fn add(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
20+
other.with_borrowed_ptr(py, |other| unsafe {
21+
err::result_from_owned_ptr(py, ffi::PyNumber_Add(self.as_ptr(), other))
22+
})
23+
}
24+
/// Perform subtraction (self - other)
25+
///
26+
/// Invokes the `__sub__` magic-method
27+
#[inline]
28+
fn subtract(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
29+
other.with_borrowed_ptr(py, |other| unsafe {
30+
err::result_from_owned_ptr(py, ffi::PyNumber_Subtract(self.as_ptr(), other))
31+
})
32+
}
33+
/// Perform multiplication (self * other)
34+
///
35+
/// Invokes the `__mul__` magic-method
36+
#[inline]
37+
fn multiply(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
38+
other.with_borrowed_ptr(py, |other| unsafe {
39+
err::result_from_owned_ptr(py, ffi::PyNumber_Multiply(self.as_ptr(), other))
40+
})
41+
}
42+
/// Perform matrix multiplication, equivalent to the Python expression `self @ other`
43+
///
44+
/// Invokes the `__matmul__` magic-method
45+
///
46+
/// This was added in Python 3.5, and will unconditionally
47+
/// throw an exception on any version before that.
48+
///
49+
/// See [PEP 0456](https://www.python.org/dev/peps/pep-0465/) for details.
50+
#[inline]
51+
fn matrix_multiply(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
52+
#[cfg(not(any(feature = "python3-4", feature = "python2-sys")))] {
53+
other.with_borrowed_ptr(py, |other| unsafe {
54+
err::result_from_owned_ptr(py, ffi::PyNumber_MatrixMultiply(self.as_ptr(), other))
55+
})
56+
}
57+
#[cfg(any(feature = "python3-4", feature = "python2-sys"))] {
58+
59+
drop(other);
60+
Err(crate::PyErr::new::<crate::exc::TypeError, _>(
61+
py,
62+
"Matrix multiplication is unsupported before Python 3.5"
63+
))
64+
}
65+
}
66+
/// Perform exponentiation, equivalent to the Python expression `self ** other`,
67+
/// or the two-argument form of the builtin method pow: `pow(self, other)`
68+
///
69+
/// Invokes the `__pow__` magic-method
70+
///
71+
/// See also [NumberProtocol::power_modulo].
72+
#[inline]
73+
fn power(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
74+
self.power_modulo(py, other, py.None())
75+
}
76+
/// Perform exponentiation modulo an integer,
77+
/// mathematically equivalent to `self ** other % mod`
78+
/// but computed much more efficiently.
79+
///
80+
/// Equivalent to invoking the three-argument form
81+
/// of the builtin `pow` method: `pow(self, other, z)`
82+
///
83+
/// Invoking with a `None` for modulo is equivalent to
84+
/// the regular power operation.
85+
///
86+
/// Invokes the `__pow__` magic-method
87+
#[inline]
88+
fn power_modulo(&self, py: Python, exp: impl ToPyObject, z: impl ToPyObject) -> PyResult<PyObject> {
89+
exp.with_borrowed_ptr(py, |exp| {
90+
z.with_borrowed_ptr(py, |z| unsafe {
91+
err::result_from_owned_ptr(py, ffi::PyNumber_Power(self.as_ptr(), exp, z))
92+
})
93+
})
94+
}
95+
/// Perform the "true division" operation,
96+
/// equivalent to the Python expression `self / other`,
97+
///
98+
/// Invokes the `__truediv__` magic-method.
99+
#[inline]
100+
fn true_divide(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
101+
other.with_borrowed_ptr(py, |other| unsafe {
102+
err::result_from_owned_ptr(py, ffi::PyNumber_TrueDivide(self.as_ptr(), other))
103+
})
104+
}
105+
/// Perform the "floor division" operation,
106+
/// equivalent to the Python expression `self // other`,
107+
///
108+
/// This method was added in Python 3.
109+
/// If compiling against Python 2, it unconditional throws an error.
110+
///
111+
/// Invokes the `__floordiv__` magic-method.
112+
#[inline]
113+
fn floor_divide(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
114+
other.with_borrowed_ptr(py, |other| unsafe {
115+
err::result_from_owned_ptr(py, ffi::PyNumber_FloorDivide(self.as_ptr(), other))
116+
})
117+
}
118+
/// Return the remainder of dividing `self` by `other`,
119+
/// equivalent to the Python expression `self % other`
120+
///
121+
/// Invokes the `__mod__` magic-method.
122+
#[inline]
123+
fn modulo(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
124+
other.with_borrowed_ptr(py, |other| unsafe {
125+
err::result_from_owned_ptr(py, ffi::PyNumber_Remainder(self.as_ptr(), other))
126+
})
127+
}
128+
/// Perform combined division and modulo,
129+
/// equivalent to the builtin method `divmod(self, other)`
130+
///
131+
/// Invokes the `__divmod__` magic-method.
132+
#[inline]
133+
fn div_mod(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
134+
other.with_borrowed_ptr(py, |other| unsafe {
135+
err::result_from_owned_ptr(py, ffi::PyNumber_Divmod(self.as_ptr(), other))
136+
})
137+
}
138+
/// Perform the negation of self (-self)
139+
///
140+
/// Invokes the `__neg__` magic-method.
141+
#[inline]
142+
fn negative(&self, py: Python) -> PyResult<PyObject> {
143+
unsafe {
144+
err::result_from_owned_ptr(py, ffi::PyNumber_Negative(self.as_ptr()))
145+
}
146+
}
147+
/// Invoke the 'positive' operation, equivalent to the
148+
/// Python expression `+self`
149+
///
150+
/// Invokes the `__pos__` magic-method
151+
#[inline]
152+
fn positive(&self, py: Python) -> PyResult<PyObject> {
153+
unsafe {
154+
err::result_from_owned_ptr(py, ffi::PyNumber_Positive(self.as_ptr()))
155+
}
156+
}
157+
/// Return the absolute value of self,
158+
/// equivalent to calling the builtin function `abs`
159+
///
160+
/// Invokes the `__abs__` magic-method.
161+
#[inline]
162+
fn absolute(&self, py: Python) -> PyResult<PyObject> {
163+
unsafe {
164+
err::result_from_owned_ptr(py, ffi::PyNumber_Absolute(self.as_ptr()))
165+
}
166+
}
167+
/// Perform the bitwise negation of self,
168+
/// equivalent to the Python expression `~self`
169+
///
170+
/// Invokes the `__invert__` magic-method
171+
#[inline]
172+
fn bitwise_invert(&self, py: Python) -> PyResult<PyObject> {
173+
unsafe {
174+
err::result_from_owned_ptr(py, ffi::PyNumber_Invert(self.as_ptr()))
175+
}
176+
}
177+
/// Shift this value to the left by the specified number of bits,
178+
/// equivalent to the Python expression `self << bits`
179+
///
180+
/// Invokes the `__lshift__` magic-method
181+
#[inline]
182+
fn left_shift(&self, py: Python, bits: impl ToPyObject) -> PyResult<PyObject> {
183+
bits.with_borrowed_ptr(py, |other| unsafe {
184+
err::result_from_owned_ptr(py, ffi::PyNumber_Lshift(self.as_ptr(), other))
185+
})
186+
}
187+
/// Shift this value to the right by the specified number of bits,
188+
/// equivalent to the Python expression `self >> bits`
189+
///
190+
/// Invokes the `__rshift__` magic-method
191+
#[inline]
192+
fn right_shift(&self, py: Python, bits: impl ToPyObject) -> PyResult<PyObject> {
193+
bits.with_borrowed_ptr(py, |other| unsafe {
194+
err::result_from_owned_ptr(py, ffi::PyNumber_Rshift(self.as_ptr(), other))
195+
})
196+
}
197+
/// Perform the "bitwise and" of `self & other`
198+
///
199+
/// Invokes the `__and__` magic-method.
200+
#[inline]
201+
fn bitwise_and(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
202+
other.with_borrowed_ptr(py, |other| unsafe {
203+
err::result_from_owned_ptr(py, ffi::PyNumber_And(self.as_ptr(), other))
204+
})
205+
}
206+
/// Perform the "bitwise exclusive or",
207+
/// equivalent to Python expression `self ^ other`
208+
///
209+
/// Invokes the `__xor__` magic-method.
210+
#[inline]
211+
fn bitwise_xor(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
212+
other.with_borrowed_ptr(py, |other| unsafe {
213+
err::result_from_owned_ptr(py, ffi::PyNumber_Xor(self.as_ptr(), other))
214+
})
215+
}
216+
/// Perform the "bitwise or" of `self | other`
217+
///
218+
/// Invokes the `__or__` magic-method.
219+
#[inline]
220+
fn bitwise_or(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
221+
other.with_borrowed_ptr(py, |other| unsafe {
222+
err::result_from_owned_ptr(py, ffi::PyNumber_Or(self.as_ptr(), other))
223+
})
224+
}
225+
/// Convert this object to an integer,
226+
/// equivalent to the builtin function `int(self)`
227+
///
228+
/// Invokes the `__int__` magic-method.
229+
///
230+
/// Throws an exception if unable to perform
231+
/// the conversion.
232+
#[inline]
233+
fn as_int(&self, py: Python) -> PyResult<PyLong> {
234+
let obj = unsafe {
235+
err::result_from_owned_ptr(py, ffi::PyNumber_Long(self.as_ptr()))?
236+
};
237+
Ok(obj.cast_into::<PyLong>(py)?)
238+
}
239+
/// Convert this object to a float,
240+
/// equivalent to the builtin function `float(self)`
241+
///
242+
/// Invokes the `__float__` magic-method.
243+
///
244+
/// Throws an exception if unable to perform
245+
/// the conversion.
246+
#[inline]
247+
fn as_float(&self, py: Python) -> PyResult<PyFloat> {
248+
let obj = unsafe {
249+
err::result_from_owned_ptr(py, ffi::PyNumber_Float(self.as_ptr()))?
250+
};
251+
Ok(obj.cast_into::<PyFloat>(py)?)
252+
}
253+
/// Losslessly convert this object to an integer index,
254+
/// as if calling `operator.index()`
255+
///
256+
/// The presence of this method indicates
257+
/// this object is an integer type.
258+
///
259+
/// Calls the `__index__` magic-method.
260+
///
261+
/// See also: [Documentation on the corresponding magic-method](https://docs.python.org/3/reference/datamodel.html?highlight=__index__#object.__index__)
262+
#[inline]
263+
fn as_int_index(&self, py: Python) -> PyResult<PyLong> {
264+
let obj = unsafe {
265+
err::result_from_owned_ptr(py, ffi::PyNumber_Index(self.as_ptr()))?
266+
};
267+
Ok(obj.cast_into::<PyLong>(py)?)
268+
}
269+
}
270+
271+
impl NumberProtocol for PyObject {}
272+
273+
274+
#[cfg(test)]
275+
mod test {
276+
use crate::*;
277+
use super::*;
278+
279+
#[test]
280+
fn addition() {
281+
let guard = Python::acquire_gil();
282+
let py = guard.python();
283+
let i1 = (5i32).to_py_object(py).into_object();
284+
let i2 = (12i32).to_py_object(py).into_object();
285+
let actual_res = i1.add(py, i2).unwrap();
286+
let expected_res = (17i32).to_py_object(py).into_object();
287+
assert_eq!(
288+
actual_res.compare(py, expected_res).unwrap(),
289+
Ordering::Equal
290+
);
291+
}
292+
293+
py_class!(class DummyMatMul |py| {
294+
data number: i32;
295+
def __new__(_cls, arg: i32) -> PyResult<DummyMatMul> {
296+
DummyMatMul::create_instance(py, arg)
297+
}
298+
def __matmul__(left, other) -> PyResult<PyObject> {
299+
// Do a dummy operation that can be easily tested
300+
left.cast_as::<Self>(py)?.number(py)
301+
.to_py_object(py)
302+
.into_object()
303+
.multiply(py, other)?
304+
.add(py, 3)
305+
}
306+
});
307+
308+
#[test]
309+
#[cfg_attr(any(feature = "python3-4", feature = "python2-sys"), should_panic)]
310+
fn matrix_multiply() {
311+
let guard = Python::acquire_gil();
312+
let py = guard.python();
313+
let first = DummyMatMul::create_instance(py, 5).unwrap().into_object();
314+
let seven = (7i32).to_py_object(py).into_object();
315+
let actual_res = first.matrix_multiply(py, seven).unwrap();
316+
// 5 * 7 + 3 => 38
317+
let expected_res = (38i32).to_py_object(py).into_object();
318+
assert_eq!(
319+
actual_res.compare(py, expected_res).unwrap(),
320+
Ordering::Equal
321+
);
322+
}
323+
}

0 commit comments

Comments
 (0)