Add as_str and godot_name to nonbitfield enums

SomeRanDev · SomeRanDev · commit 44f38dbe7a6d · 2024-09-18T11:59:40.000-04:00
diff --git a/godot-codegen/src/generator/enums.rs b/godot-codegen/src/generator/enums.rs
@@ -76,7 +76,7 @@ pub fn make_enum_definition_with(
                 quote! { #[doc(hidden)] pub }
             });
 
-            let debug_impl = make_enum_debug_impl(enum_);
+            let debug_impl = make_enum_debug_impl(enum_, define_traits && !enum_.is_bitfield);
             quote! {
                 #[repr(transparent)]
                 #[derive( #( #derives ),* )]
@@ -146,11 +146,8 @@ fn make_enum_index_impl(enum_: &Enum) -> Option<TokenStream> {
     })
 }
 
-/// Implement `Debug` trait for the enum.
-fn make_enum_debug_impl(enum_: &Enum) -> TokenStream {
-    let enum_name = &enum_.name;
-    let enum_name_str = enum_name.to_string();
-
+// Creates the match cases to return the enumerator name as &str.
+fn make_enum_to_str_cases(enum_: &Enum) -> TokenStream {
     let enumerators = enum_.enumerators.iter().map(|enumerator| {
         let Enumerator { name, .. } = enumerator;
         let name_str = name.to_string();
@@ -159,22 +156,55 @@ fn make_enum_debug_impl(enum_: &Enum) -> TokenStream {
         }
     });
 
+    quote! {
+        #( #enumerators )*
+    }
+}
+
+/// Implement `Debug` trait for the enum.
+fn make_enum_debug_impl(enum_: &Enum, use_as_str: bool) -> TokenStream {
+    let enum_name = &enum_.name;
+    let enum_name_str = enum_name.to_string();
+
+    // Print the ord if no matching enumerator can be found.
+    let enumerator_not_found = quote! {
+        f.debug_struct(#enum_name_str)
+            .field("ord", &self.ord)
+            .finish()?;
+
+        return Ok(());
+    };
+
+    // Reuse `as_str` if traits are defined and not a bitfield.
+    let function_body = if use_as_str {
+        quote! {
+            use crate::obj::EngineEnum;
+
+            let enumerator = self.as_str();
+            if enumerator.is_empty() {
+                #enumerator_not_found
+            }
+        }
+    } else {
+        let enumerators = make_enum_to_str_cases(enum_);
+
+        quote! {
+            // Many enums have duplicates, thus allow unreachable.
+            // In the future, we could print sth like "ONE|TWO" instead (at least for unstable Debug).
+            #[allow(unreachable_patterns)]
+            let enumerator = match *self {
+                #enumerators
+                _ => {
+                    #enumerator_not_found
+                }
+            };
+        }
+    };
+
     quote! {
         impl std::fmt::Debug for #enum_name {
             fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-                // Many enums have duplicates, thus allow unreachable.
-                // In the future, we could print sth like "ONE|TWO" instead (at least for unstable Debug).
-                #[allow(unreachable_patterns)]
-                let enumerator = match *self {
-                    #( #enumerators )*
-                    _ => {
-                        f.debug_struct(#enum_name_str)
-                        .field("ord", &self.ord)
-                        .finish()?;
-                        return Ok(());
-                    }
-                };
-
+                #function_body
                 f.write_str(enumerator)
             }
         }
@@ -222,6 +252,8 @@ fn make_enum_engine_trait_impl(enum_: &Enum) -> TokenStream {
             }
         });
 
+        let str_functions = make_enum_str_functions(enum_);
+
         quote! {
             impl #engine_trait for #name {
                 fn try_from_ord(ord: i32) -> Option<Self> {
@@ -234,10 +266,13 @@ fn make_enum_engine_trait_impl(enum_: &Enum) -> TokenStream {
                 fn ord(self) -> i32 {
                     self as i32
                 }
+
+                #str_functions
             }
         }
     } else {
         let unique_ords = enum_.unique_ords().expect("self is an enum");
+        let str_functions = make_enum_str_functions(enum_);
 
         quote! {
             impl #engine_trait for #name {
@@ -251,8 +286,70 @@ fn make_enum_engine_trait_impl(enum_: &Enum) -> TokenStream {
                 fn ord(self) -> i32 {
                     self.ord
                 }
+
+                #str_functions
+            }
+        }
+    }
+}
+
+/// Creates the `as_str` and `godot_name` implementations for the enum.
+fn make_enum_str_functions(enum_: &Enum) -> TokenStream {
+    let as_str_enumerators = make_enum_to_str_cases(enum_);
+
+    // Only enumerations with different godot names are specified.
+    // `as_str` is called for the rest of them.
+    let godot_different_cases = {
+        let enumerators = enum_
+            .enumerators
+            .iter()
+            .filter(|enumerator| enumerator.name != enumerator.godot_name)
+            .map(|enumerator| {
+                let Enumerator {
+                    name, godot_name, ..
+                } = enumerator;
+                let godot_name_str = godot_name.to_string();
+                quote! {
+                    Self::#name => #godot_name_str,
+                }
+            });
+
+        quote! {
+            #( #enumerators )*
+        }
+    };
+
+    let godot_name_match = if godot_different_cases.is_empty() {
+        // If empty, all the Rust names match the Godot ones.
+        // Remove match statement to avoid `clippy::match_single_binding`.
+        quote! {
+            self.as_str()
+        }
+    } else {
+        quote! {
+            // Many enums have duplicates, thus allow unreachable.
+            #[allow(unreachable_patterns)]
+            match *self {
+                #godot_different_cases
+                _ => self.as_str(),
             }
         }
+    };
+
+    quote! {
+        #[inline]
+        fn as_str(&self) -> &'static str {
+            // Many enums have duplicates, thus allow unreachable.
+            #[allow(unreachable_patterns)]
+            match *self {
+                #as_str_enumerators
+                _ => "",
+            }
+        }
+
+        fn godot_name(&self) -> &'static str {
+            #godot_name_match
+        }
     }
 }
 
diff --git a/godot-core/src/builtin/vectors/vector_axis.rs b/godot-core/src/builtin/vectors/vector_axis.rs
@@ -38,6 +38,22 @@ macro_rules! impl_vector_axis_enum {
             fn ord(self) -> i32 {
                 self as i32
             }
+
+            fn as_str(&self) -> &'static str {
+                match *self {
+                    $(
+                        Self::$axis => stringify!($axis),
+                    )+
+                }
+            }
+
+            fn godot_name(&self) -> &'static str {
+                match *self {
+                    $(
+                        Self::$axis => concat!("AXIS_", stringify!($axis)),
+                    )+
+                }
+            }
         }
 
         impl GodotConvert for $AxisEnum {
diff --git a/godot-core/src/obj/traits.rs b/godot-core/src/obj/traits.rs
@@ -164,6 +164,16 @@ pub trait EngineEnum: Copy {
         Self::try_from_ord(ord)
             .unwrap_or_else(|| panic!("ordinal {ord} does not map to any enumerator"))
     }
+
+    // The name of the enumerator, as it appears in Rust.
+    //
+    // If the value does not match one of the known enumerators, the empty string is returned.
+    fn as_str(&self) -> &'static str;
+
+    // The equivalent name of the enumerator, as specified in Godot.
+    //
+    // If the value does not match one of the known enumerators, the empty string is returned.
+    fn godot_name(&self) -> &'static str;
 }
 
 /// Auto-implemented for all engine-provided bitfields.
diff --git a/itest/rust/src/object_tests/enum_test.rs b/itest/rust/src/object_tests/enum_test.rs
@@ -10,6 +10,7 @@ use godot::builtin::varray;
 use godot::classes::input::CursorShape;
 use godot::classes::mesh::PrimitiveType;
 use godot::classes::{time, ArrayMesh};
+use godot::global::{Orientation, Key};
 use std::collections::HashSet;
 
 #[itest]
@@ -70,3 +71,29 @@ fn add_surface_from_arrays() {
     let mut mesh = ArrayMesh::new();
     mesh.add_surface_from_arrays(PrimitiveType::PRIMITIVE_TRIANGLES, varray![]);
 }
+
+#[itest]
+fn enum_as_str_correct() {
+    use godot::obj::EngineEnum;
+    assert_eq!(Orientation::Vertical.as_str(), "VERTICAL");
+    assert_eq!(Orientation::Horizontal.as_str(), "HORIZONTAL");
+
+    assert_eq!(Key::NONE.as_str(), "NONE");
+    assert_eq!(Key::SPECIAL.as_str(), "SPECIAL");
+    assert_eq!(Key::ESCAPE.as_str(), "ESCAPE");
+    assert_eq!(Key::TAB.as_str(), "TAB");
+    assert_eq!(Key::A.as_str(), "A");
+}
+
+#[itest]
+fn enum_godot_name_correct() {
+    use godot::obj::EngineEnum;
+    assert_eq!(Orientation::Vertical.godot_name(), Orientation::Vertical.as_str());
+    assert_eq!(Orientation::Horizontal.godot_name(), Orientation::Vertical.as_str());
+
+    assert_eq!(Key::NONE.godot_name(), "KEY_NONE");
+    assert_eq!(Key::SPECIAL.godot_name(), "KEY_SPECIAL");
+    assert_eq!(Key::ESCAPE.godot_name(), "KEY_ESCAPE");
+    assert_eq!(Key::TAB.godot_name(), "KEY_TAB");
+    assert_eq!(Key::A.godot_name(), "KEY_A");
+}
