refactor(stackable-versioned): Rework version enum for merged_crd

stackabletech · Sep 24, 2024 · 8253c2b · 8253c2b
1 parent 34c6c4f
commit 8253c2b
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 35 deletions.
diff --git a/crates/stackable-versioned-macros/src/codegen/vstruct/mod.rs b/crates/stackable-versioned-macros/src/codegen/vstruct/mod.rs
@@ -319,22 +319,40 @@ impl VersionedStruct {
         crd_fn_calls: Vec<TokenStream>,
         enum_variants: Vec<(Ident, String)>,
     ) -> TokenStream {
-        let ident = &self.idents.kubernetes;
+        let enum_ident = &self.idents.kubernetes;
+        let enum_vis = &self.visibility;
 
-        let version_enum_definition = self.generate_kubernetes_version_enum(enum_variants);
+        let mut enum_variant_matches = TokenStream::new();
+        let mut enum_variant_idents = TokenStream::new();
+
+        for (enum_variant_ident, enum_variant_display) in enum_variants {
+            enum_variant_idents.extend(quote! {#enum_variant_ident,});
+            enum_variant_matches.extend(quote! {
+                #enum_ident::#enum_variant_ident => f.write_str(#enum_variant_display),
+            });
+        }
 
         quote! {
             #[automatically_derived]
-            pub struct #ident;
+            #enum_vis enum #enum_ident {
+                #enum_variant_idents
+            }
 
-            #version_enum_definition
+            #[automatically_derived]
+            impl ::std::fmt::Display for #enum_ident {
+                fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::result::Result<(), ::std::fmt::Error> {
+                    match self {
+                        #enum_variant_matches
+                    }
+                }
+            }
 
             #[automatically_derived]
-            impl #ident {
+            impl #enum_ident {
                 /// Generates a merged CRD which contains all versions defined using the
                 /// `#[versioned()]` macro.
                 pub fn merged_crd(
-                    stored_apiversion: Version
+                    stored_apiversion: Self
                 ) -> ::std::result::Result<::k8s_openapi::apiextensions_apiserver::pkg::apis::apiextensions::v1::CustomResourceDefinition, ::kube::core::crd::MergeError> {
                     ::kube::core::crd::merge_crds(vec![#(#crd_fn_calls),*], &stored_apiversion.to_string())
                 }
@@ -353,32 +371,4 @@ impl VersionedStruct {
             <#path as ::kube::CustomResourceExt>::crd()
         }
     }
-
-    fn generate_kubernetes_version_enum(&self, enum_variants: Vec<(Ident, String)>) -> TokenStream {
-        let mut enum_variant_matches = TokenStream::new();
-        let mut enum_variant_idents = TokenStream::new();
-
-        for (enum_variant_ident, enum_variant_display) in enum_variants {
-            enum_variant_idents.extend(quote! {#enum_variant_ident,});
-            enum_variant_matches.extend(quote! {
-                Version::#enum_variant_ident => f.write_str(#enum_variant_display),
-            });
-        }
-
-        quote! {
-            #[automatically_derived]
-            pub enum Version {
-                #enum_variant_idents
-            }
-
-            #[automatically_derived]
-            impl ::std::fmt::Display for Version {
-                fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::result::Result<(), ::std::fmt::Error> {
-                    match self {
-                        #enum_variant_matches
-                    }
-                }
-            }
-        }
-    }
 }
diff --git a/crates/stackable-versioned-macros/tests/k8s/pass/crd.rs b/crates/stackable-versioned-macros/tests/k8s/pass/crd.rs
@@ -26,6 +26,6 @@ fn main() {
         baz: bool,
     }
 
-    let merged_crd = Foo::merged_crd(Version::V1).unwrap();
+    let merged_crd = Foo::merged_crd(Foo::V1).unwrap();
     println!("{}", serde_yaml::to_string(&merged_crd).unwrap());
 }