Implement default slot methods and default __repr__ for enums.

Author: jovenlin0527

pyo3-macros-backend/src/pyclass.rs

@@ -3,7 +3,7 @@
 use crate::attributes::{self, take_pyo3_options, NameAttribute, TextSignatureAttribute};
 use crate::deprecations::Deprecations;
 use crate::konst::{ConstAttributes, ConstSpec};
-use crate::pyimpl::{gen_py_const, PyClassMethodsType};
+use crate::pyimpl::{gen_default_slot_impls, gen_py_const, PyClassMethodsType};
 use crate::pymethod::{impl_py_getter_def, impl_py_setter_def, PropertyType};
 use crate::utils::{self, unwrap_group, PythonDoc};
 use proc_macro2::{Span, TokenStream};
@@ -422,6 +422,23 @@ fn impl_enum_class(
         .impl_all();
     let descriptors = unit_variants_as_descriptors(cls, variants.iter().map(|v| v.ident));
 
+    let variants_repr = variants.iter().map(|variant| {
+        let variant_name = variant.ident;
+        // Assuming all variants are unit variants because they are the only type we support.
+        let repr = format!("{}.{}", cls, variant_name);
+        quote! { #cls::#variant_name => #repr }
+    });
+
+    let default_repr_impl = quote! {
+        fn __repr__(&self) -> &'static str {
+            match self {
+                #(#variants_repr),*
+                , // needs an additional comma for the last item
+                _ => unreachable!("Unsupported variant type."),
+            }
+        }
+    };
+    let default_impls = gen_default_slot_impls(cls, vec![default_repr_impl]);
     Ok(quote! {
 
         #pytypeinfo
@@ -430,6 +447,8 @@ fn impl_enum_class(
 
         #descriptors
 
+        #default_impls
+
     })
 }
 

pyo3-macros-backend/src/pyimpl.rs

@@ -8,7 +8,7 @@ use crate::{
     pymethod::{self, is_proto_method},
 };
 use proc_macro2::TokenStream;
-use pymethod::GeneratedPyMethod;
+use pymethod::{GeneratedPyMethod, PyMethodKind, PyMethodProtoKind};
 use quote::quote;
 use syn::spanned::Spanned;
 
@@ -139,6 +139,59 @@ pub fn gen_py_const(cls: &syn::Type, spec: &ConstSpec) -> TokenStream {
     }
 }
 
+pub fn gen_default_slot_impls(cls: &syn::Ident, method_defs: Vec<TokenStream>) -> TokenStream {
+    // This function uses a lot of `unwrap()`; since method_defs are provided by us, they should
+    // all succeed.
+    let ty: syn::Type = syn::parse_quote!(#cls);
+
+    let mut method_defs: Vec<syn::ImplItemMethod> = method_defs
+        .into_iter()
+        .map(|tokens| syn::parse2(tokens).unwrap())
+        .collect();
+
+    // Generate the python method definition, and change the rust method name
+    // to avoid name collision.
+    let py_methods: Vec<_> = method_defs
+        .iter_mut()
+        .map(|method_def| -> TokenStream {
+            // The new name of internal method.
+            let new_name: syn::Ident =
+                syn::parse_str(&format!("__pyo3__{}", method_def.sig.ident.clone())).unwrap();
+            let mut method = pymethod::PyMethod::parse(
+                &mut method_def.sig,
+                &mut method_def.attrs,
+                Default::default(),
+            )
+            .unwrap();
+            // Have to change the name after parsing so the slot is parsed correctly.
+            method.spec.name = &new_name;
+            let tok = match method.kind {
+                PyMethodKind::Proto(proto_kind) => match proto_kind {
+                    PyMethodProtoKind::Slot(slot_def) => {
+                        slot_def.generate_type_slot(&ty, &method.spec).unwrap()
+                    }
+                    _ => todo!("Unsupported prototype kind"),
+                },
+                _ => unreachable!("Not a prototype method"),
+            };
+            // Change the name here because `method` mutably borrows from `method_def`
+            method_def.sig.ident = new_name;
+            tok
+        })
+        .collect(); // So `py_methods` don't mutably borrow from `method_defs`
+    quote! {
+        impl #cls {
+            #(#method_defs)*
+        }
+        impl ::pyo3::class::impl_::PyClassDefaultSlots<#ty>
+            for ::pyo3::class::impl_::PyClassImplCollector<#ty> {
+                fn py_class_default_slots(self) -> &'static [::pyo3::ffi::PyType_Slot] {
+                    &[#(#py_methods),*]
+                }
+        }
+    }
+}
+
 fn impl_py_methods(ty: &syn::Type, methods: Vec<TokenStream>) -> TokenStream {
     quote! {
         impl ::pyo3::class::impl_::PyMethods<#ty>

pyo3-macros-backend/src/pymethod.rs

@@ -25,12 +25,12 @@ pub enum GeneratedPyMethod {
 }
 
 pub struct PyMethod<'a> {
-    kind: PyMethodKind,
+    pub kind: PyMethodKind,
     method_name: String,
-    spec: FnSpec<'a>,
+    pub spec: FnSpec<'a>,
 }
 
-enum PyMethodKind {
+pub enum PyMethodKind {
     Fn,
     Proto(PyMethodProtoKind),
 }
@@ -49,14 +49,14 @@ impl PyMethodKind {
     }
 }
 
-enum PyMethodProtoKind {
+pub enum PyMethodProtoKind {
     Slot(&'static SlotDef),
     Call,
     SlotFragment(&'static SlotFragmentDef),
 }
 
 impl<'a> PyMethod<'a> {
-    fn parse(
+    pub fn parse(
         sig: &'a mut syn::Signature,
         meth_attrs: &mut Vec<syn::Attribute>,
         options: PyFunctionOptions,
@@ -744,7 +744,7 @@ impl ReturnMode {
     }
 }
 
-struct SlotDef {
+pub struct SlotDef {
     slot: StaticIdent,
     func_ty: StaticIdent,
     arguments: &'static [Ty],
@@ -799,7 +799,7 @@ impl SlotDef {
         self
     }
 
-    fn generate_type_slot(&self, cls: &syn::Type, spec: &FnSpec) -> Result<TokenStream> {
+    pub fn generate_type_slot(&self, cls: &syn::Type, spec: &FnSpec) -> Result<TokenStream> {
         let SlotDef {
             slot,
             func_ty,
@@ -871,7 +871,7 @@ fn generate_method_body(
     })
 }
 
-struct SlotFragmentDef {
+pub struct SlotFragmentDef {
     fragment: &'static str,
     arguments: &'static [Ty],
     extract_error_mode: ExtractErrorMode,

tests/test_default_impls.rs

@@ -1,118 +1,42 @@
 #![allow(non_snake_case)]
-use pyo3::class::PyMethodDefType;
 use pyo3::prelude::*;
-use pyo3::py_run;
 
 mod common;
 
-// Tests for PyClassDefaultSlots
+// Test default generated __repr__.
 #[pyclass]
-struct TestDefaultSlots;
-
-// generated using `Cargo expand`
-// equivalent to
-// ```
-// impl TestDefaultSlots {{
-//     fn __str__(&self) -> &'static str {
-//         "default"
-//     }
-// }
-// ```
-impl TestDefaultSlots {
-    fn __pyo3__str__(&self) -> &'static str {
-        "default"
-    }
-}
-
-impl ::pyo3::class::impl_::PyClassDefaultSlots<TestDefaultSlots>
-    for ::pyo3::class::impl_::PyClassImplCollector<TestDefaultSlots>
-{
-    fn py_class_default_slots(self) -> &'static [::pyo3::ffi::PyType_Slot] {
-        &[{
-            unsafe extern "C" fn __wrap(
-                _raw_slf: *mut ::pyo3::ffi::PyObject,
-            ) -> *mut ::pyo3::ffi::PyObject {
-                let _slf = _raw_slf;
-                ::pyo3::callback::handle_panic(|_py| {
-                    let _cell = _py
-                        .from_borrowed_ptr::<::pyo3::PyAny>(_slf)
-                        .downcast::<::pyo3::PyCell<TestDefaultSlots>>()?;
-                    let _ref = _cell.try_borrow()?;
-                    let _slf = &_ref;
-                    ::pyo3::callback::convert(_py, TestDefaultSlots::__pyo3__str__(_slf))
-                })
-            }
-            ::pyo3::ffi::PyType_Slot {
-                slot: ::pyo3::ffi::Py_tp_str,
-                pfunc: __wrap as ::pyo3::ffi::reprfunc as _,
-            }
-        }]
-    }
+enum TestDefaultRepr {
+    Var,
 }
 
 #[test]
 fn test_default_slot_exists() {
     Python::with_gil(|py| {
-        let test_object = Py::new(py, TestDefaultSlots).unwrap();
-        py_assert!(py, test_object, "str(test_object) == 'default'");
+        let test_object = Py::new(py, TestDefaultRepr::Var).unwrap();
+        py_assert!(
+            py,
+            test_object,
+            "repr(test_object) == 'TestDefaultRepr.Var'"
+        );
     })
 }
 
 #[pyclass]
-struct OverrideSlot;
-
-// generated using `Cargo expand`
-// equivalent to
-// ```
-// impl OverrideMagicMethod {
-//     fn __str__(&self) -> &'static str {
-//         "default"
-//     }
-// }
-// ```
-impl OverrideSlot {
-    fn __pyo3__str__(&self) -> &'static str {
-        "default"
-    }
-}
-
-impl ::pyo3::class::impl_::PyClassDefaultSlots<OverrideSlot>
-    for ::pyo3::class::impl_::PyClassImplCollector<OverrideSlot>
-{
-    fn py_class_default_slots(self) -> &'static [::pyo3::ffi::PyType_Slot] {
-        &[{
-            unsafe extern "C" fn __wrap(
-                _raw_slf: *mut ::pyo3::ffi::PyObject,
-            ) -> *mut ::pyo3::ffi::PyObject {
-                let _slf = _raw_slf;
-                ::pyo3::callback::handle_panic(|_py| {
-                    let _cell = _py
-                        .from_borrowed_ptr::<::pyo3::PyAny>(_slf)
-                        .downcast::<::pyo3::PyCell<OverrideSlot>>()?;
-                    let _ref = _cell.try_borrow()?;
-                    let _slf = &_ref;
-                    ::pyo3::callback::convert(_py, OverrideSlot::__pyo3__str__(_slf))
-                })
-            }
-            ::pyo3::ffi::PyType_Slot {
-                slot: ::pyo3::ffi::Py_tp_str,
-                pfunc: __wrap as ::pyo3::ffi::reprfunc as _,
-            }
-        }]
-    }
+enum OverrideSlot {
+    Var,
 }
 
 #[pymethods]
 impl OverrideSlot {
-    fn __str__(&self) -> &str {
+    fn __repr__(&self) -> &str {
         "overriden"
     }
 }
 
 #[test]
 fn test_override_slot() {
     Python::with_gil(|py| {
-        let test_object = Py::new(py, OverrideSlot).unwrap();
-        py_assert!(py, test_object, "str(test_object) == 'overriden'");
+        let test_object = Py::new(py, OverrideSlot::Var).unwrap();
+        py_assert!(py, test_object, "repr(test_object) == 'overriden'");
     })
 }