Refactor PyMethodDef creation too.

PyO3 · Jun 5, 2021 · 1878b9d · 1878b9d
1 parent 53f9370
commit 1878b9d
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 49 deletions.
diff --git a/pyo3-macros-backend/src/method.rs b/pyo3-macros-backend/src/method.rs
@@ -572,6 +572,35 @@ impl<'a> FnSpec<'a> {
             }
         })
     }
+
+    pub fn get_methoddef(&self, wrapper: impl ToTokens) -> TokenStream {
+        let python_name = self.null_terminated_python_name();
+        let doc = &self.doc;
+        match self.convention {
+            CallingConvention::Noargs => quote! {
+                pyo3::class::methods::PyMethodDef::noargs(
+                    #python_name,
+                    pyo3::class::methods::PyCFunction(#wrapper),
+                    #doc,
+                )
+            },
+            CallingConvention::Fastcall => quote! {
+                pyo3::class::methods::PyMethodDef::fastcall_cfunction_with_keywords(
+                    #python_name,
+                    pyo3::class::methods::PyCFunctionFastWithKeywords(#wrapper),
+                    #doc,
+                )
+            },
+            CallingConvention::Varargs => quote! {
+                pyo3::class::methods::PyMethodDef::cfunction_with_keywords(
+                    #python_name,
+                    pyo3::class::methods::PyCFunctionWithKeywords(#wrapper),
+                    #doc,
+                )
+            },
+            CallingConvention::TpNew => unreachable!("tp_new cannot get a methoddef"),
+        }
+    }
 }
 
 #[derive(Clone, PartialEq, Debug)]

diff --git a/pyo3-macros-backend/src/pyfunction.rs b/pyo3-macros-backend/src/pyfunction.rs
@@ -400,37 +400,16 @@ pub fn impl_wrap_pyfunction(
         deprecations: options.deprecations,
     };
 
-    let doc = &spec.doc;
-    let python_name = spec.null_terminated_python_name();
-
-    let name = &func.sig.ident;
-    let wrapper_ident = format_ident!("__pyo3_raw_{}", name);
+    let wrapper_ident = format_ident!("__pyo3_raw_{}", spec.name);
     let wrapper = spec.get_wrapper_function(&wrapper_ident, None)?;
-    let (methoddef_meth, cfunc_variant) = match spec.convention {
-        CallingConvention::Noargs => (quote!(noargs), quote!(PyCFunction)),
-        CallingConvention::Fastcall => (
-            quote!(fastcall_cfunction_with_keywords),
-            quote!(PyCFunctionFastWithKeywords),
-        ),
-        _ => (
-            quote!(cfunction_with_keywords),
-            quote!(PyCFunctionWithKeywords),
-        ),
-    };
+    let methoddef = spec.get_methoddef(wrapper_ident);
 
     let wrapped_pyfunction = quote! {
         #wrapper
         pub(crate) fn #function_wrapper_ident<'a>(
             args: impl Into<pyo3::derive_utils::PyFunctionArguments<'a>>
         ) -> pyo3::PyResult<&'a pyo3::types::PyCFunction> {
-            pyo3::types::PyCFunction::internal_new(
-                pyo3::class::methods::PyMethodDef:: #methoddef_meth (
-                    #python_name,
-                    pyo3::class::methods:: #cfunc_variant (#wrapper_ident),
-                    #doc,
-                ),
-                args.into(),
-            )
+            pyo3::types::PyCFunction::internal_new(#methoddef, args.into())
         }
     };
     Ok((function_wrapper_ident, wrapped_pyfunction))

diff --git a/pyo3-macros-backend/src/pymethod.rs b/pyo3-macros-backend/src/pymethod.rs
@@ -4,7 +4,6 @@ use std::borrow::Cow;
 
 use crate::attributes::NameAttribute;
 use crate::konst::ConstSpec;
-use crate::method::CallingConvention;
 use crate::utils::ensure_not_async_fn;
 use crate::{deprecations::Deprecations, utils};
 use crate::{
@@ -242,36 +241,14 @@ pub fn impl_py_method_def(
     let wrapper_ident = syn::Ident::new("__wrap", Span::call_site());
     let wrapper_def = spec.get_wrapper_function(&wrapper_ident, Some(cls))?;
     let add_flags = flags.map(|flags| quote!(.flags(#flags)));
-    let doc = &spec.doc;
-    let python_name = spec.null_terminated_python_name();
     let methoddef_type = match spec.tp {
         FnType::FnStatic => quote!(Static),
         FnType::FnClass => quote!(Class),
         _ => quote!(Method),
     };
-    let (methoddef_meth, cfunc_variant) = match spec.convention {
-        CallingConvention::Noargs => (quote!(noargs), quote!(PyCFunction)),
-        CallingConvention::Fastcall => (
-            quote!(fastcall_cfunction_with_keywords),
-            quote!(PyCFunctionFastWithKeywords),
-        ),
-        _ => (
-            quote!(cfunction_with_keywords),
-            quote!(PyCFunctionWithKeywords),
-        ),
-    };
+    let methoddef = spec.get_methoddef(quote! {{ #wrapper_def #wrapper_ident }});
     Ok(quote! {
-        pyo3::class::PyMethodDefType:: #methoddef_type ({
-            pyo3::class::PyMethodDef:: #methoddef_meth (
-                #python_name,
-                pyo3::class::methods:: #cfunc_variant ({
-                    #wrapper_def
-                    #wrapper_ident
-                }),
-                #doc
-            )
-            #add_flags
-        })
+        pyo3::class::PyMethodDefType::#methoddef_type(#methoddef #add_flags)
     })
 }