feat: add enum const types

dojoengine · Mar 10, 2025 · 55dc93c · 55dc93c
1 parent 5b83295
commit 55dc93c
Show file tree

Hide file tree

Showing 2 changed files with 133 additions and 17 deletions.
diff --git a/crates/dojo/bindgen/src/plugins/typescript/generator/enum.rs b/crates/dojo/bindgen/src/plugins/typescript/generator/enum.rs
@@ -1,7 +1,8 @@
 use cainome::parser::tokens::{Composite, CompositeType};
+use convert_case::{Case, Casing};
 
 use super::constants::{CAIRO_ENUM_IMPORT, CAIRO_ENUM_TOKEN, SN_IMPORT_SEARCH};
-use super::token_is_enum;
+use super::{token_is_custom_enum, token_is_enum};
 use crate::error::BindgenResult;
 use crate::plugins::typescript::generator::JsPrimitiveType;
 use crate::plugins::{BindgenModelGenerator, Buffer};
@@ -21,33 +22,80 @@ impl TsEnumGenerator {
             }
         }
     }
-}
 
-impl BindgenModelGenerator for TsEnumGenerator {
-    fn generate(&self, token: &Composite, buffer: &mut Buffer) -> BindgenResult<String> {
-        if token.r#type != CompositeType::Enum || token.inners.is_empty() {
-            return Ok(String::new());
-        }
+    fn generate_simple_enum(
+        &self,
+        token: &Composite,
+        buffer: &mut Buffer,
+    ) -> BindgenResult<String> {
+        Ok(format!(
+            "// Type definition for `{path}` enum
+export const {camel_name} = [
+{variants}
+] as const;
+export type {name} = {{ [key in typeof {camel_name}[number]]: string }};
+export type {name}Enum = CairoCustomEnum;
+",
+            path = token.type_path,
+            name = token.type_name(),
+            camel_name = token.type_name().to_case(Case::Camel),
+            variants = token
+                .inners
+                .iter()
+                .map(|inner| { format!("\t'{}',", inner.name) })
+                .collect::<Vec<String>>()
+                .join("\n")
+        ))
+    }
 
-        self.check_import(token, buffer);
-        let gen = format!(
+    fn generate_custom_enum(
+        &self,
+        token: &Composite,
+        buffer: &mut Buffer,
+    ) -> BindgenResult<String> {
+        Ok(format!(
             "// Type definition for `{path}` enum
-export type {name} = {{
+export const {camel_name} = [
 {variants}
-}}
+] as const;
+export type {name} = {{ 
+{variant_types}
+}};
 export type {name}Enum = CairoCustomEnum;
 ",
             path = token.type_path,
             name = token.type_name(),
+            camel_name = token.type_name().to_case(Case::Camel),
             variants = token
+                .inners
+                .iter()
+                .map(|inner| { format!("\t'{}',", inner.name) })
+                .collect::<Vec<String>>()
+                .join("\n"),
+            variant_types = token
                 .inners
                 .iter()
                 .map(|inner| {
-                    format!("\t{}: {};", inner.name, JsPrimitiveType::from(&inner.token))
+                    format!("\t{}: {},", inner.name, JsPrimitiveType::from(&inner.token))
                 })
                 .collect::<Vec<String>>()
                 .join("\n")
-        );
+        ))
+    }
+}
+
+impl BindgenModelGenerator for TsEnumGenerator {
+    fn generate(&self, token: &Composite, buffer: &mut Buffer) -> BindgenResult<String> {
+        if token.r#type != CompositeType::Enum || token.inners.is_empty() {
+            return Ok(String::new());
+        }
+
+        self.check_import(token, buffer);
+        let gen = if token_is_custom_enum(token) {
+            self.generate_custom_enum(token, buffer)?
+        } else {
+            self.generate_simple_enum(token, buffer)?
+        };
 
         if buffer.has(&gen) {
             return Ok(String::new());
@@ -107,8 +155,9 @@ mod tests {
 
         assert_eq!(
             result,
-            "// Type definition for `core::test::AvailableTheme` enum\nexport type AvailableTheme \
-             = {\n\tLight: string;\n\tDark: string;\n\tDojo: string;\n}\nexport type \
+            "// Type definition for `core::test::AvailableTheme` enum\nexport const \
+             availableTheme = [\n\t'Light',\n\t'Dark',\n\t'Dojo',\n] as const;\nexport type \
+             AvailableTheme = { [key in typeof availableTheme[number]]: string };\nexport type \
              AvailableThemeEnum = CairoCustomEnum;\n"
         );
     }
@@ -118,8 +167,9 @@ mod tests {
         let mut buff = Buffer::new();
         let writer = TsEnumGenerator;
         buff.push(
-            "// Type definition for `core::test::AvailableTheme` enum\nexport type AvailableTheme \
-             = {\n\tLight: string;\n\tDark: string;\n\tDojo: string;\n}\nexport type \
+            "// Type definition for `core::test::AvailableTheme` enum\nexport const \
+             availableTheme = [\n\t'Light',\n\t'Dark',\n\t'Dojo',\n] as const;\nexport type \
+             AvailableTheme = { [key in typeof availableTheme[number]]: string };\nexport type \
              AvailableThemeEnum = CairoCustomEnum;\n"
                 .to_owned(),
         );
@@ -131,6 +181,21 @@ mod tests {
         assert!(result.is_empty())
     }
 
+    #[test]
+    fn test_custom_enum() {
+        let mut buff = Buffer::new();
+        let writer = TsEnumGenerator;
+        let token = create_custom_enum_token();
+        let result = writer.generate(&token, &mut buff).unwrap();
+        assert_eq!(
+            result,
+            "// Type definition for `core::test::CustomEnum` enum\nexport const customEnum = \
+             [\n\t'Predefined',\n\t'Custom',\n] as const;\nexport type CustomEnum = { \
+             \n\tPredefined: AvailableThemeEnum,\n\tCustom: Custom,\n};\nexport type \
+             CustomEnumEnum = CairoCustomEnum;\n"
+        );
+    }
+
     fn create_available_theme_enum_token() -> Composite {
         Composite {
             type_path: "core::test::AvailableTheme".to_owned(),
@@ -160,4 +225,39 @@ mod tests {
             alias: None,
         }
     }
+    fn create_custom_enum_token() -> Composite {
+        Composite {
+            type_path: "core::test::CustomEnum".to_owned(),
+            inners: vec![
+                CompositeInner {
+                    index: 0,
+                    name: "Predefined".to_owned(),
+                    kind: CompositeInnerKind::Key,
+                    token: Token::Composite(create_available_theme_enum_token()),
+                },
+                CompositeInner {
+                    index: 1,
+                    name: "Custom".to_owned(),
+                    kind: CompositeInnerKind::Key,
+                    token: Token::Composite(Composite {
+                        type_path: "core::test::custom::Custom".to_owned(),
+                        inners: vec![CompositeInner {
+                            index: 0,
+                            name: "Classname".to_owned(),
+                            kind: CompositeInnerKind::Key,
+                            token: Token::CoreBasic(CoreBasic { type_path: "felt252".to_owned() }),
+                        }],
+                        generic_args: vec![],
+                        r#type: CompositeType::Struct,
+                        is_event: false,
+                        alias: None,
+                    }),
+                },
+            ],
+            generic_args: vec![],
+            r#type: CompositeType::Enum,
+            is_event: false,
+            alias: None,
+        }
+    }
 }
diff --git a/crates/dojo/bindgen/src/plugins/typescript/generator/mod.rs b/crates/dojo/bindgen/src/plugins/typescript/generator/mod.rs
@@ -48,13 +48,29 @@ pub(crate) fn generate_type_init(token: &Composite) -> String {
 pub(crate) fn token_is_option(token: &Composite) -> bool {
     token.type_path.starts_with(CAIRO_OPTION_TYPE_PATH)
 }
+/// Checks if token has inner composite
+/// * token - The token to check
+fn token_has_inner_composite(token: &Composite) -> bool {
+    token.inners.iter().any(|inner| match &inner.token {
+        Token::Array(array) => array.inner.to_composite().is_ok(),
+        Token::Tuple(tuple) => tuple.inners.iter().any(|t| matches!(t, Token::Composite(_))),
+        Token::Composite(_) => true,
+        _ => false,
+    })
+}
 
 /// Checks if Token::Composite is an custom enum (enum with nested Composite types)
 /// * token - The token to check
 pub(crate) fn token_is_enum(token: &Composite) -> bool {
     token.r#type == CompositeType::Enum
 }
 
+/// Checks if Token::Composite is an custom enum (enum with nested Composite types)
+/// * token - The token to check
+pub(crate) fn token_is_custom_enum(token: &Composite) -> bool {
+    token.r#type == CompositeType::Enum && token_has_inner_composite(token)
+}
+
 /// Type used to map cainome `Token` into javascript types in interface definition
 #[derive(Debug)]
 pub(crate) struct JsPrimitiveType(String);