customize into_py() for enum PyClass

pyo3-macros-backend/src/pyclass.rs

@@ -791,7 +791,7 @@ fn impl_complex_enum(
 
     let default_slots = vec![];
 
-    let pyclass_impls = PyClassImplsBuilder::new(
+    let impl_builder = PyClassImplsBuilder::new(
         cls,
         &args,
         methods_type,
@@ -801,8 +801,45 @@ fn impl_complex_enum(
         ),
         default_slots,
     )
-    .doc(doc)
-    .impl_all()?;
+    .doc(doc);
+
+    // Need to customize the into_py impl so that it returns the variant PyClass
+    let enum_into_py_impl = {
+        let match_arms: Vec<TokenStream> = variants
+            .iter()
+            .map(|variant| {
+                let variant_ident = variant.ident;
+                let variant_cls = gen_complex_enum_variant_class_ident(cls, variant.ident);
+                quote! {
+                    #cls::#variant_ident { .. } => {
+                        let pyclass_init = _pyo3::PyClassInitializer::from(self).add_subclass(#variant_cls);
+                        let variant_value = _pyo3::Py::new(py, pyclass_init).unwrap();
+                        _pyo3::IntoPy::into_py(variant_value, py)
+                    }
+                }
+            })
+            .collect();
+
+        quote! {
+            impl _pyo3::IntoPy<_pyo3::PyObject> for #cls {
+                fn into_py(self, py: _pyo3::Python) -> _pyo3::PyObject {
+                    match self {
+                        #(#match_arms)*
+                    }
+                }
+            }
+        }
+    };
+
+    let pyclass_impls: TokenStream = vec![
+        impl_builder.impl_pyclass(),
+        impl_builder.impl_extractext(),
+        enum_into_py_impl,
+        impl_builder.impl_pyclassimpl()?,
+        impl_builder.impl_freelist(),
+    ]
+    .into_iter()
+    .collect();
 
     let mut variant_cls_zsts = vec![];
     let mut variant_cls_pytypeinfos = vec![];