Merge pull request #791 from kngwyu/pyclass-arg

Make it enable to take &PyClass as arguments as pyfunctions/methods
PyO3 · Mar 4, 2020 · 74b22eb · 74b22eb
2 parents c5a093d + 96115ea
commit 74b22eb
Show file tree

Hide file tree

Showing 8 changed files with 195 additions and 150 deletions.
diff --git a/pyo3-derive-backend/src/method.rs b/pyo3-derive-backend/src/method.rs
@@ -109,12 +109,11 @@ impl<'a> FnSpec<'a> {
 
                     let py = crate::utils::if_type_is_python(ty);
 
-                    let opt = check_arg_ty_and_optional(name, ty);
+                    let opt = check_ty_optional(ty);
                     arguments.push(FnArg {
                         name: ident,
                         by_ref,
                         mutability,
-                        // mode: mode,
                         ty,
                         optional: opt,
                         py,
@@ -305,55 +304,18 @@ pub fn is_ref(name: &syn::Ident, ty: &syn::Type) -> bool {
     false
 }
 
-pub fn check_arg_ty_and_optional<'a>(
-    name: &'a syn::Ident,
-    ty: &'a syn::Type,
-) -> Option<&'a syn::Type> {
-    match ty {
-        syn::Type::Path(syn::TypePath { ref path, .. }) => {
-            //if let Some(ref qs) = qs {
-            //    panic!("explicit Self type in a 'qualified path' is not supported: {:?} - {:?}",
-            //           name, qs);
-            //}
-
-            if let Some(segment) = path.segments.last() {
-                match segment.ident.to_string().as_str() {
-                    "Option" => match segment.arguments {
-                        syn::PathArguments::AngleBracketed(ref params) => {
-                            if params.args.len() != 1 {
-                                panic!("argument type is not supported by python method: {:?} ({:?}) {:?}",
-                                       name,
-                                       ty,
-                                       path);
-                            }
-
-                            match &params.args[0] {
-                                syn::GenericArgument::Type(ref ty) => Some(ty),
-                                _ => panic!("argument type is not supported by python method: {:?} ({:?}) {:?}",
-                                            name,
-                                            ty,
-                                            path),
-                            }
-                        }
-                        _ => {
-                            panic!(
-                                "argument type is not supported by python method: {:?} ({:?}) {:?}",
-                                name, ty, path
-                            );
-                        }
-                    },
-                    _ => None,
-                }
-            } else {
-                None
-            }
-        }
-        _ => {
-            None
-            //panic!("argument type is not supported by python method: {:?} ({:?})",
-            //name,
-            //ty);
-        }
+pub(crate) fn check_ty_optional<'a>(ty: &'a syn::Type) -> Option<&'a syn::Type> {
+    let path = match ty {
+        syn::Type::Path(syn::TypePath { ref path, .. }) => path,
+        _ => return None,
+    };
+    let seg = path.segments.last().filter(|s| s.ident == "Option")?;
+    match seg.arguments {
+        syn::PathArguments::AngleBracketed(ref params) => match params.args.first() {
+            Some(syn::GenericArgument::Type(ref ty)) => Some(ty),
+            _ => None,
+        },
+        _ => None,
     }
 }
 

diff --git a/pyo3-derive-backend/src/module.rs b/pyo3-derive-backend/src/module.rs
@@ -68,7 +68,7 @@ fn wrap_fn_argument<'a>(cap: &'a syn::PatType, name: &'a Ident) -> syn::Result<m
     };
 
     let py = crate::utils::if_type_is_python(&cap.ty);
-    let opt = method::check_arg_ty_and_optional(&name, &cap.ty);
+    let opt = method::check_ty_optional(&cap.ty);
     Ok(method::FnArg {
         name: ident,
         mutability,

diff --git a/pyo3-derive-backend/src/pyclass.rs b/pyo3-derive-backend/src/pyclass.rs
@@ -412,6 +412,15 @@ fn impl_class(
             type BaseNativeType = #base_nativetype;
         }
 
+        impl<'a> pyo3::derive_utils::ExtractExt<'a> for &'a #cls
+        {
+            type Target = pyo3::PyRef<'a, #cls>;
+        }
+        impl<'a> pyo3::derive_utils::ExtractExt<'a> for &'a mut #cls
+        {
+            type Target = pyo3::PyRefMut<'a, #cls>;
+        }
+
         #into_pyobject
 
         #inventory_impl

diff --git a/pyo3-derive-backend/src/pymethod.rs b/pyo3-derive-backend/src/pymethod.rs
@@ -387,10 +387,7 @@ pub(crate) fn impl_wrap_setter(
             };
             match _result {
                 Ok(_) => 0,
-                Err(e) => {
-                    e.restore(_py);
-                    -1
-                }
+                Err(e) => e.restore_and_minus1(_py),
             }
         }
     })
@@ -523,45 +520,65 @@ fn impl_arg_param(
     }
     let arg_value = quote!(output[#option_pos]);
     *option_pos += 1;
-    if arg.optional.is_some() {
-        let default = if let Some(d) = spec.default_value(name) {
-            if d.to_string() == "None" {
-                quote! { None }
-            } else {
-                quote! { Some(#d) }
-            }
+
+    return if let Some(ty) = arg.optional.as_ref() {
+        let default = if let Some(d) = spec.default_value(name).filter(|d| d.to_string() != "None")
+        {
+            quote! { Some(#d) }
         } else {
             quote! { None }
         };
-        quote! {
-            let #arg_name = match #arg_value.as_ref() {
-                Some(_obj) => {
-                    if _obj.is_none() {
-                        #default
-                    } else {
-                        Some(_obj.extract()?)
-                    }
-                },
-                None => #default
+        if let syn::Type::Reference(tref) = ty {
+            let (tref, mut_) = tref_preprocess(tref);
+            let as_deref = if mut_.is_some() {
+                quote! { as_deref_mut }
+            } else {
+                quote! { as_deref }
             };
+            // Get Option<&T> from Option<PyRef<T>>
+            quote! {
+                let #mut_ _tmp = match #arg_value.as_ref().filter(|obj| !obj.is_none()) {
+                    Some(_obj) => {
+                        Some(_obj.extract::<<#tref as pyo3::derive_utils::ExtractExt>::Target>()?)
+                    },
+                    None => #default,
+                };
+                let #arg_name = _tmp.#as_deref();
+            }
+        } else {
+            quote! {
+                let #arg_name = match #arg_value.as_ref().filter(|obj| !obj.is_none()) {
+                    Some(_obj) => Some(_obj.extract()?),
+                    None => #default,
+                };
+            }
         }
     } else if let Some(default) = spec.default_value(name) {
         quote! {
-            let #arg_name = match #arg_value.as_ref() {
-                Some(_obj) => {
-                    if _obj.is_none() {
-                        #default
-                    } else {
-                        _obj.extract()?
-                    }
-                },
-                None => #default
+            let #arg_name = match #arg_value.as_ref().filter(|obj| !obj.is_none()) {
+                Some(_obj) => _obj.extract()?,
+                None => #default,
             };
         }
+    } else if let syn::Type::Reference(tref) = arg.ty {
+        let (tref, mut_) = tref_preprocess(tref);
+        // Get &T from PyRef<T>
+        quote! {
+            let #mut_ _tmp: <#tref as pyo3::derive_utils::ExtractExt>::Target
+                = #arg_value.unwrap().extract()?;
+            let #arg_name = &#mut_ *_tmp;
+        }
     } else {
         quote! {
             let #arg_name = #arg_value.unwrap().extract()?;
         }
+    };
+
+    fn tref_preprocess(tref: &syn::TypeReference) -> (syn::TypeReference, Option<syn::token::Mut>) {
+        let mut tref = tref.to_owned();
+        tref.lifetime = None;
+        let mut_ = tref.mutability;
+        (tref, mut_)
     }
 }
 

diff --git a/src/derive_utils.rs b/src/derive_utils.rs
@@ -198,6 +198,7 @@ impl<T: PyClass, I: Into<PyClassInitializer<T>>> IntoPyNewResult<T, I> for PyRes
     }
 }
 
+#[doc(hidden)]
 pub trait GetPropertyValue {
     fn get_property_value(&self, py: Python) -> PyObject;
 }
@@ -218,6 +219,7 @@ impl GetPropertyValue for PyObject {
 }
 
 /// Utilities for basetype
+#[doc(hidden)]
 pub trait PyBaseTypeUtils {
     type Dict;
     type WeakRef;
@@ -231,3 +233,16 @@ impl<T: PyClass> PyBaseTypeUtils for T {
     type LayoutAsBase = crate::pycell::PyCellInner<T>;
     type BaseNativeType = T::BaseNativeType;
 }
+
+/// Utility trait to enable &PyClass as a pymethod/function argument
+#[doc(hidden)]
+pub trait ExtractExt<'a> {
+    type Target: crate::FromPyObject<'a>;
+}
+
+impl<'a, T> ExtractExt<'a> for T
+where
+    T: crate::FromPyObject<'a>,
+{
+    type Target = T;
+}
diff --git a/tests/test_dunder.rs b/tests/test_dunder.rs