Merge pull request #2083 from aviramha/magic_methods

verify py method args count
PyO3 · Jan 7, 2022 · 2cee7fe · 2cee7fe
2 parents e01add5 + cbfb5ac
commit 2cee7fe
Show file tree

Hide file tree

Showing 16 changed files with 520 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,6 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Expose `pyo3-build-config` APIs for cross-compiling and Python configuration discovery for use in other projects. [#1996](https://github.com/PyO3/pyo3/pull/1996)
 - Add buffer magic methods `__getbuffer__` and `__releasebuffer__` to `#[pymethods]`. [#2067](https://github.com/PyO3/pyo3/pull/2067)
 - Accept paths in `wrap_pyfunction` and `wrap_pymodule`. [#2081](https://github.com/PyO3/pyo3/pull/2081)
+- Add check for correct number of arguments on magic methods. [#2083](https://github.com/PyO3/pyo3/pull/2083)
 
 ### Changed
 
@@ -46,6 +47,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - internal `handle_panic` helper [#2074](https://github.com/PyO3/pyo3/pull/2074)
   - `#[pyfunction]` and `#[pymethods]` argument extraction [#2075](https://github.com/PyO3/pyo3/pull/2075) [#2085](https://github.com/PyO3/pyo3/pull/2085)
   - `#[pyclass]` type object creation [#2076](https://github.com/PyO3/pyo3/pull/2076) [#2081](https://github.com/PyO3/pyo3/pull/2081)
+- `__ipow__` now supports modulo argument on Python 3.8+. [#2083](https://github.com/PyO3/pyo3/pull/2083)
+- `pyo3-macros-backend` is now compiled with PyO3 cfgs to enable different magic method definitions based on version. [#2083](https://github.com/PyO3/pyo3/pull/2083)
 
 ### Removed
 

diff --git a/guide/src/class/protocols.md b/guide/src/class/protocols.md
@@ -289,13 +289,15 @@ This trait also has support the augmented arithmetic assignments (`+=`, `-=`,
   * `fn __itruediv__(&'p mut self, other: impl FromPyObject) -> PyResult<()>`
   * `fn __ifloordiv__(&'p mut self, other: impl FromPyObject) -> PyResult<()>`
   * `fn __imod__(&'p mut self, other: impl FromPyObject) -> PyResult<()>`
-  * `fn __ipow__(&'p mut self, other: impl FromPyObject) -> PyResult<()>`
+  * `fn __ipow__(&'p mut self, other: impl FromPyObject, modulo: impl FromPyObject) -> PyResult<()>` on Python 3.8^
+  * `fn __ipow__(&'p mut self, other: impl FromPyObject) -> PyResult<()>` on Python 3.7 see https://bugs.python.org/issue36379
   * `fn __ilshift__(&'p mut self, other: impl FromPyObject) -> PyResult<()>`
   * `fn __irshift__(&'p mut self, other: impl FromPyObject) -> PyResult<()>`
   * `fn __iand__(&'p mut self, other: impl FromPyObject) -> PyResult<()>`
   * `fn __ior__(&'p mut self, other: impl FromPyObject) -> PyResult<()>`
   * `fn __ixor__(&'p mut self, other: impl FromPyObject) -> PyResult<()>`
 
+
 The following methods implement the unary arithmetic operations (`-`, `+`, `abs()` and `~`):
 
   * `fn __neg__(&'p self) -> PyResult<impl ToPyObject>`

diff --git a/pyo3-macros-backend/src/defs.rs b/pyo3-macros-backend/src/defs.rs
@@ -421,7 +421,7 @@ pub const NUM: Proto = Proto {
             .args(&["Other"])
             .has_self(),
         MethodProto::new("__ipow__", "PyNumberIPowProtocol")
-            .args(&["Other"])
+            .args(&["Other", "Modulo"])
             .has_self(),
         MethodProto::new("__ilshift__", "PyNumberILShiftProtocol")
             .args(&["Other"])

diff --git a/pyo3-macros-backend/src/pymethod.rs b/pyo3-macros-backend/src/pymethod.rs
@@ -532,8 +532,8 @@ const __IMOD__: SlotDef = SlotDef::new("Py_nb_inplace_remainder", "binaryfunc")
     .arguments(&[Ty::Object])
     .extract_error_mode(ExtractErrorMode::NotImplemented)
     .return_self();
-const __IPOW__: SlotDef = SlotDef::new("Py_nb_inplace_power", "ternaryfunc")
-    .arguments(&[Ty::Object, Ty::Object])
+const __IPOW__: SlotDef = SlotDef::new("Py_nb_inplace_power", "ipowfunc")
+    .arguments(&[Ty::Object, Ty::IPowModulo])
     .extract_error_mode(ExtractErrorMode::NotImplemented)
     .return_self();
 const __ILSHIFT__: SlotDef = SlotDef::new("Py_nb_inplace_lshift", "binaryfunc")
@@ -613,6 +613,7 @@ enum Ty {
     Object,
     MaybeNullObject,
     NonNullObject,
+    IPowModulo,
     CompareOp,
     Int,
     PyHashT,
@@ -626,6 +627,7 @@ impl Ty {
         match self {
             Ty::Object | Ty::MaybeNullObject => quote! { *mut _pyo3::ffi::PyObject },
             Ty::NonNullObject => quote! { ::std::ptr::NonNull<_pyo3::ffi::PyObject> },
+            Ty::IPowModulo => quote! { _pyo3::impl_::pymethods::IPowModulo },
             Ty::Int | Ty::CompareOp => quote! { ::std::os::raw::c_int },
             Ty::PyHashT => quote! { _pyo3::ffi::Py_hash_t },
             Ty::PySsizeT => quote! { _pyo3::ffi::Py_ssize_t },
@@ -679,6 +681,16 @@ impl Ty {
                 );
                 extract_object(cls, arg.ty, ident, extract)
             }
+            Ty::IPowModulo => {
+                let extract = handle_error(
+                    extract_error_mode,
+                    py,
+                    quote! {
+                        #ident.extract(#py)
+                    },
+                );
+                extract_object(cls, arg.ty, ident, extract)
+            }
             Ty::CompareOp => {
                 let extract = handle_error(
                     extract_error_mode,
@@ -888,8 +900,11 @@ fn generate_method_body(
 ) -> Result<TokenStream> {
     let self_conversion = spec.tp.self_conversion(Some(cls), extract_error_mode);
     let rust_name = spec.name;
-    let (arg_idents, conversions) =
+    let (arg_idents, arg_count, conversions) =
         extract_proto_arguments(cls, py, &spec.args, arguments, extract_error_mode)?;
+    if arg_count != arguments.len() {
+        bail_spanned!(spec.name.span() => format!("Expected {} arguments, got {}", arguments.len(), arg_count));
+    }
     let call = quote! { _pyo3::callback::convert(#py, #cls::#rust_name(_slf, #(#arg_idents),*)) };
     let body = if let Some(return_mode) = return_mode {
         return_mode.return_call_output(py, call)
@@ -1058,7 +1073,7 @@ fn extract_proto_arguments(
     method_args: &[FnArg],
     proto_args: &[Ty],
     extract_error_mode: ExtractErrorMode,
-) -> Result<(Vec<Ident>, TokenStream)> {
+) -> Result<(Vec<Ident>, usize, TokenStream)> {
     let mut arg_idents = Vec::with_capacity(method_args.len());
     let mut non_python_args = 0;
 
@@ -1077,9 +1092,8 @@ fn extract_proto_arguments(
             arg_idents.push(ident);
         }
     }
-
     let conversions = quote!(#(#args_conversions)*);
-    Ok((arg_idents, conversions))
+    Ok((arg_idents, non_python_args, conversions))
 }
 
 struct StaticIdent(&'static str);

diff --git a/src/class/number.rs b/src/class/number.rs
@@ -221,7 +221,7 @@ pub trait PyNumberProtocol<'p>: PyClass {
     {
         unimplemented!()
     }
-    fn __ipow__(&'p mut self, other: Self::Other) -> Self::Result
+    fn __ipow__(&'p mut self, other: Self::Other, modulo: Option<Self::Modulo>) -> Self::Result
     where
         Self: PyNumberIPowProtocol<'p>,
     {
@@ -504,6 +504,8 @@ pub trait PyNumberIDivmodProtocol<'p>: PyNumberProtocol<'p> {
 pub trait PyNumberIPowProtocol<'p>: PyNumberProtocol<'p> {
     type Other: FromPyObject<'p>;
     type Result: IntoPyCallbackOutput<()>;
+    // See https://bugs.python.org/issue36379
+    type Modulo: FromPyObject<'p>;
 }
 
 #[allow(clippy::upper_case_acronyms)]
@@ -718,17 +720,28 @@ py_binary_self_func!(imod, PyNumberIModProtocol, T::__imod__);
 pub unsafe extern "C" fn ipow<T>(
     slf: *mut ffi::PyObject,
     other: *mut ffi::PyObject,
-    _modulo: *mut ffi::PyObject,
+    modulo: crate::impl_::pymethods::IPowModulo,
 ) -> *mut ffi::PyObject
 where
     T: for<'p> PyNumberIPowProtocol<'p>,
 {
-    // NOTE: Somehow __ipow__ causes SIGSEGV in Python < 3.8 when we extract,
-    // so we ignore it. It's the same as what CPython does.
     crate::callback_body!(py, {
         let slf_cell = py.from_borrowed_ptr::<crate::PyCell<T>>(slf);
         let other = py.from_borrowed_ptr::<crate::PyAny>(other);
-        call_operator_mut!(py, slf_cell, __ipow__, other).convert(py)?;
+        slf_cell
+            .try_borrow_mut()?
+            .__ipow__(
+                extract_or_return_not_implemented!(other),
+                match modulo.extract(py) {
+                    Ok(value) => value,
+                    Err(_) => {
+                        let res = crate::ffi::Py_NotImplemented();
+                        crate::ffi::Py_INCREF(res);
+                        return Ok(res);
+                    }
+                },
+            )
+            .convert(py)?;
         ffi::Py_INCREF(slf);
         Ok::<_, PyErr>(slf)
     })

diff --git a/src/ffi/mod.rs b/src/ffi/mod.rs
@@ -204,3 +204,7 @@ mod cpython;
 
 #[cfg(not(Py_LIMITED_API))]
 pub use self::cpython::*;
+
+/// Helper to enable #\[pymethods\] to see the workaround for __ipow__ on Python 3.7
+#[doc(hidden)]
+pub use crate::impl_::pymethods::ipowfunc;
diff --git a/src/impl_.rs b/src/impl_.rs
@@ -12,4 +12,6 @@ pub mod frompyobject;
 pub(crate) mod not_send;
 #[doc(hidden)]
 pub mod pyclass;
+#[doc(hidden)]
+pub mod pymethods;
 pub mod pymodule;
diff --git a/src/impl_/pymethods.rs b/src/impl_/pymethods.rs
@@ -0,0 +1,33 @@
+use crate::{ffi, FromPyObject, PyAny, PyResult, Python};
+
+/// Python 3.8 and up - __ipow__ has modulo argument correctly populated.
+#[cfg(Py_3_8)]
+#[repr(transparent)]
+pub struct IPowModulo(*mut ffi::PyObject);
+
+/// Python 3.7 and older - __ipow__ does not have modulo argument correctly populated.
+#[cfg(not(Py_3_8))]
+#[repr(transparent)]
+pub struct IPowModulo(std::mem::MaybeUninit<*mut ffi::PyObject>);
+
+/// Helper to use as pymethod ffi definition
+#[allow(non_camel_case_types)]
+pub type ipowfunc = unsafe extern "C" fn(
+    arg1: *mut ffi::PyObject,
+    arg2: *mut ffi::PyObject,
+    arg3: IPowModulo,
+) -> *mut ffi::PyObject;
+
+impl IPowModulo {
+    #[cfg(Py_3_8)]
+    #[inline]
+    pub fn extract<'a, T: FromPyObject<'a>>(self, py: Python<'a>) -> PyResult<T> {
+        unsafe { py.from_borrowed_ptr::<PyAny>(self.0) }.extract()
+    }
+
+    #[cfg(not(Py_3_8))]
+    #[inline]
+    pub fn extract<'a, T: FromPyObject<'a>>(self, py: Python<'a>) -> PyResult<T> {
+        unsafe { py.from_borrowed_ptr::<PyAny>(ffi::Py_None()) }.extract()
+    }
+}