Merge pull request #3461 from davidhewitt/some-wraps

better `Some`-wrapping for default arguments

newsfragments/3461.fixed.md

@@ -0,0 +1 @@
+Some-wrapping of `Option<T>` default arguments will no longer re-wrap `Some(T)` or expressions evaluating to `None`.

pyo3-macros-backend/src/params.rs

@@ -1,6 +1,7 @@
 use crate::{
     method::{FnArg, FnSpec},
     pyfunction::FunctionSignature,
+    quotes::some_wrap,
 };
 use proc_macro2::{Span, TokenStream};
 use quote::{quote, quote_spanned};
@@ -192,12 +193,7 @@ fn impl_arg_param(
     // Option<T> arguments have special treatment: the default should be specified _without_ the
     // Some() wrapper. Maybe this should be changed in future?!
     if arg.optional.is_some() {
-        default = Some(match &default {
-            Some(expression) if expression.to_string() != "None" => {
-                quote!(::std::option::Option::Some(#expression))
-            }
-            _ => quote!(::std::option::Option::None),
-        })
+        default = Some(default.map_or_else(|| quote!(::std::option::Option::None), some_wrap));
     }
 
     let tokens = if let Some(expr_path) = arg.attrs.from_py_with.as_ref().map(|attr| &attr.value) {

pyo3-macros-backend/src/quotes.rs

@@ -1,10 +1,16 @@
 use proc_macro2::TokenStream;
 use quote::quote;
 
+pub(crate) fn some_wrap(obj: TokenStream) -> TokenStream {
+    quote! {
+        _pyo3::impl_::wrap::SomeWrap::wrap(#obj)
+    }
+}
+
 pub(crate) fn ok_wrap(obj: TokenStream) -> TokenStream {
     quote! {
-        _pyo3::impl_::pymethods::OkWrap::wrap(#obj, py)
-            .map_err(::core::convert::Into::into)
+        _pyo3::impl_::wrap::OkWrap::wrap(#obj, py)
+            .map_err(::core::convert::Into::<_pyo3::PyErr>::into)
     }
 }
 

src/impl_.rs

@@ -19,3 +19,4 @@ pub mod pymethods;
 pub mod pymodule;
 #[doc(hidden)]
 pub mod trampoline;
+pub mod wrap;

src/impl_/pymethods.rs

@@ -1,10 +1,7 @@
 use crate::gil::LockGIL;
 use crate::impl_::panic::PanicTrap;
 use crate::internal_tricks::extract_c_string;
-use crate::{
-    ffi, IntoPy, Py, PyAny, PyCell, PyClass, PyErr, PyObject, PyResult, PyTraverseError, PyVisit,
-    Python,
-};
+use crate::{ffi, PyAny, PyCell, PyClass, PyObject, PyResult, PyTraverseError, PyVisit, Python};
 use std::borrow::Cow;
 use std::ffi::CStr;
 use std::fmt;
@@ -295,32 +292,6 @@ pub(crate) struct PyMethodDefDestructor {
     doc: Cow<'static, CStr>,
 }
 
-// The macros need to Ok-wrap the output of user defined functions; i.e. if they're not a result, make them into one.
-pub trait OkWrap<T> {
-    type Error;
-    fn wrap(self, py: Python<'_>) -> Result<Py<PyAny>, Self::Error>;
-}
-
-impl<T> OkWrap<T> for T
-where
-    T: IntoPy<PyObject>,
-{
-    type Error = PyErr;
-    fn wrap(self, py: Python<'_>) -> PyResult<Py<PyAny>> {
-        Ok(self.into_py(py))
-    }
-}
-
-impl<T, E> OkWrap<T> for Result<T, E>
-where
-    T: IntoPy<PyObject>,
-{
-    type Error = E;
-    fn wrap(self, py: Python<'_>) -> Result<Py<PyAny>, Self::Error> {
-        self.map(|o| o.into_py(py))
-    }
-}
-
 pub(crate) fn get_name(name: &'static str) -> PyResult<Cow<'static, CStr>> {
     extract_c_string(name, "function name cannot contain NUL byte.")
 }

src/impl_/wrap.rs

@@ -0,0 +1,60 @@
+use crate::{IntoPy, Py, PyAny, PyErr, PyObject, PyResult, Python};
+
+/// Used to wrap values in `Option<T>` for default arguments.
+pub trait SomeWrap<T> {
+    fn wrap(self) -> T;
+}
+
+impl<T> SomeWrap<Option<T>> for T {
+    fn wrap(self) -> Option<T> {
+        Some(self)
+    }
+}
+
+impl<T> SomeWrap<Option<T>> for Option<T> {
+    fn wrap(self) -> Self {
+        self
+    }
+}
+
+/// Used to wrap the result of `#[pyfunction]` and `#[pymethods]`.
+pub trait OkWrap<T> {
+    type Error;
+    fn wrap(self, py: Python<'_>) -> Result<Py<PyAny>, Self::Error>;
+}
+
+// The T: IntoPy<PyObject> bound here is necessary to prevent the
+// implementation for Result<T, E> from conflicting
+impl<T> OkWrap<T> for T
+where
+    T: IntoPy<PyObject>,
+{
+    type Error = PyErr;
+    fn wrap(self, py: Python<'_>) -> PyResult<Py<PyAny>> {
+        Ok(self.into_py(py))
+    }
+}
+
+impl<T, E> OkWrap<T> for Result<T, E>
+where
+    T: IntoPy<PyObject>,
+{
+    type Error = E;
+    fn wrap(self, py: Python<'_>) -> Result<Py<PyAny>, Self::Error> {
+        self.map(|o| o.into_py(py))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn wrap_option() {
+        let a: Option<u8> = SomeWrap::wrap(42);
+        assert_eq!(a, Some(42));
+
+        let b: Option<u8> = SomeWrap::wrap(None);
+        assert_eq!(b, None);
+    }
+}

tests/test_pyfunction.rs

@@ -508,3 +508,23 @@ fn test_return_value_borrows_from_arguments() {
         py_assert!(py, function key value, "function(key, value) == { \"key\": 42 }");
     });
 }
+
+#[test]
+fn test_some_wrap_arguments() {
+    // https://github.com/PyO3/pyo3/issues/3460
+    const NONE: Option<u8> = None;
+    #[pyfunction(signature = (a = 1, b = Some(2), c = None, d = NONE))]
+    fn some_wrap_arguments(
+        a: Option<u8>,
+        b: Option<u8>,
+        c: Option<u8>,
+        d: Option<u8>,
+    ) -> [Option<u8>; 4] {
+        [a, b, c, d]
+    }
+
+    Python::with_gil(|py| {
+        let function = wrap_pyfunction!(some_wrap_arguments, py).unwrap();
+        py_assert!(py, function, "function() == [1, 2, None, None]");
+    })
+}