feat(codegen): implemented oneOf unions for primitive types

torn-api-codegen/src/model/enum.rs

@@ -1,12 +1,15 @@
-use heck::ToUpperCamelCase;
+use heck::{ToSnakeCase, ToUpperCamelCase};
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
+use syn::Ident;
 
 use crate::openapi::{
     parameter::OpenApiParameterSchema,
     r#type::{OpenApiType, OpenApiVariants},
 };
 
+use super::{object::PrimitiveType, ResolvedSchema};
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum EnumRepr {
     U8,
@@ -17,10 +20,12 @@ pub enum EnumRepr {
 pub enum EnumVariantTupleValue {
     Ref { ty_name: String },
     ArrayOfRefs { ty_name: String },
+    Primitive(PrimitiveType),
+    Enum { name: String, inner: Enum },
 }
 
 impl EnumVariantTupleValue {
-    pub fn from_schema(schema: &OpenApiType) -> Option<Self> {
+    pub fn from_schema(name: &str, schema: &OpenApiType) -> Option<Self> {
         match schema {
             OpenApiType {
                 ref_path: Some(path),
@@ -44,14 +49,61 @@ impl EnumVariantTupleValue {
                     ty_name: path.strip_prefix("#/components/schemas/")?.to_owned(),
                 })
             }
+            OpenApiType {
+                r#type: Some("string"),
+                format: None,
+                r#enum: None,
+                ..
+            } => Some(Self::Primitive(PrimitiveType::String)),
+            OpenApiType {
+                r#type: Some("string"),
+                format: None,
+                r#enum: Some(_),
+                ..
+            } => {
+                let name = format!("{name}Variant");
+                Some(Self::Enum {
+                    inner: Enum::from_schema(&name, schema)?,
+                    name,
+                })
+            }
+            OpenApiType {
+                r#type: Some("integer"),
+                format: Some("int64"),
+                ..
+            } => Some(Self::Primitive(PrimitiveType::I64)),
+            OpenApiType {
+                r#type: Some("integer"),
+                format: Some("int32"),
+                ..
+            } => Some(Self::Primitive(PrimitiveType::I32)),
             _ => None,
         }
     }
 
-    pub fn type_name(&self) -> &str {
+    pub fn type_name(&self, ns: &mut EnumNamespace) -> TokenStream {
         match self {
-            Self::Ref { ty_name } => ty_name,
-            Self::ArrayOfRefs { ty_name } => ty_name,
+            Self::Ref { ty_name } => {
+                let ty = format_ident!("{ty_name}");
+                quote! { crate::models::#ty }
+            }
+            Self::ArrayOfRefs { ty_name } => {
+                let ty = format_ident!("{ty_name}");
+                quote! { Vec<crate::models::#ty> }
+            }
+            Self::Primitive(PrimitiveType::I64) => quote! { i64 },
+            Self::Primitive(PrimitiveType::I32) => quote! { i32 },
+            Self::Primitive(PrimitiveType::Float) => quote! { f32 },
+            Self::Primitive(PrimitiveType::String) => quote! { String },
+            Self::Primitive(PrimitiveType::DateTime) => quote! { chrono::DateTime<chrono::Utc> },
+            Self::Primitive(PrimitiveType::Bool) => quote! { bool },
+            Self::Enum { name, .. } => {
+                let path = ns.get_ident();
+                let ty_name = format_ident!("{name}");
+                quote! {
+                    #path::#ty_name,
+                }
+            }
         }
     }
 
@@ -59,6 +111,36 @@ impl EnumVariantTupleValue {
         match self {
             Self::Ref { ty_name } => ty_name.clone(),
             Self::ArrayOfRefs { ty_name } => format!("{ty_name}s"),
+            Self::Primitive(PrimitiveType::I64) => "I64".to_owned(),
+            Self::Primitive(PrimitiveType::I32) => "I32".to_owned(),
+            Self::Primitive(PrimitiveType::Float) => "Float".to_owned(),
+            Self::Primitive(PrimitiveType::String) => "String".to_owned(),
+            Self::Primitive(PrimitiveType::DateTime) => "DateTime".to_owned(),
+            Self::Primitive(PrimitiveType::Bool) => "Bool".to_owned(),
+            Self::Enum { .. } => "Variant".to_owned(),
+        }
+    }
+
+    pub fn is_display(&self, resolved: &ResolvedSchema) -> bool {
+        match self {
+            Self::Primitive(_) => true,
+            Self::Ref { ty_name } | Self::ArrayOfRefs { ty_name } => resolved
+                .models
+                .get(ty_name)
+                .map(|f| f.is_display(resolved))
+                .unwrap_or_default(),
+            Self::Enum { inner, .. } => inner.is_display(resolved),
+        }
+    }
+
+    pub fn codegen_display(&self) -> TokenStream {
+        match self {
+            Self::ArrayOfRefs { .. } => quote! {
+                write!(f, "{}", value.iter().map(ToString::to_string).collect::<Vec<_>>().join(","))
+            },
+            _ => quote! {
+                write!(f, "{}", value)
+            },
         }
     }
 }
@@ -77,12 +159,32 @@ impl Default for EnumVariantValue {
 }
 
 impl EnumVariantValue {
-    pub fn codegen_display(&self, name: &str) -> Option<TokenStream> {
+    pub fn is_display(&self, resolved: &ResolvedSchema) -> bool {
         match self {
-            Self::Repr(i) => Some(quote! { write!(f, "{}", #i) }),
+            Self::Repr(_) | Self::String { .. } => true,
+            Self::Tuple(val) => {
+                val.len() == 1
+                    && val
+                        .iter()
+                        .next()
+                        .map(|v| v.is_display(resolved))
+                        .unwrap_or_default()
+            }
+        }
+    }
+
+    pub fn codegen_display(&self, name: &str) -> Option<TokenStream> {
+        let variant = format_ident!("{name}");
+
+        match self {
+            Self::Repr(i) => Some(quote! { Self::#variant => write!(f, "{}", #i) }),
             Self::String { rename } => {
                 let name = rename.as_deref().unwrap_or(name);
-                Some(quote! { write!(f, #name) })
+                Some(quote! { Self::#variant => write!(f, #name) })
+            }
+            Self::Tuple(values) if values.len() == 1 => {
+                let rhs = values.first().unwrap().codegen_display();
+                Some(quote! { Self::#variant(value) => #rhs })
             }
             _ => None,
         }
@@ -96,8 +198,61 @@ pub struct EnumVariant {
     pub value: EnumVariantValue,
 }
 
+pub struct EnumNamespace<'e> {
+    r#enum: &'e Enum,
+    ident: Option<Ident>,
+    elements: Vec<TokenStream>,
+    top_level_elements: Vec<TokenStream>,
+}
+
+impl EnumNamespace<'_> {
+    pub fn get_ident(&mut self) -> Ident {
+        self.ident
+            .get_or_insert_with(|| {
+                let name = self.r#enum.name.to_snake_case();
+                format_ident!("{name}")
+            })
+            .clone()
+    }
+
+    pub fn push_element(&mut self, el: TokenStream) {
+        self.elements.push(el);
+    }
+
+    pub fn push_top_level(&mut self, el: TokenStream) {
+        self.top_level_elements.push(el);
+    }
+
+    pub fn codegen(mut self) -> Option<TokenStream> {
+        if self.elements.is_empty() && self.top_level_elements.is_empty() {
+            None
+        } else {
+            let top_level = &self.top_level_elements;
+            let mut output = quote! {
+                #(#top_level)*
+            };
+
+            if !self.elements.is_empty() {
+                let ident = self.get_ident();
+                let elements = self.elements;
+                output.extend(quote! {
+                    pub mod #ident {
+                        #(#elements)*
+                    }
+                });
+            }
+
+            Some(output)
+        }
+    }
+}
+
 impl EnumVariant {
-    pub fn codegen(&self) -> Option<TokenStream> {
+    pub fn codegen(
+        &self,
+        ns: &mut EnumNamespace,
+        resolved: &ResolvedSchema,
+    ) -> Option<TokenStream> {
         let doc = self.description.as_ref().map(|d| {
             quote! {
                 #[doc = #d]
@@ -127,15 +282,29 @@ impl EnumVariant {
             EnumVariantValue::Tuple(values) => {
                 let mut val_tys = Vec::with_capacity(values.len());
 
-                for value in values {
-                    let ty_name = value.type_name();
-                    let ty_name = format_ident!("{ty_name}");
+                if let [value] = values.as_slice() {
+                    let enum_name = format_ident!("{}", ns.r#enum.name);
+                    let ty_name = value.type_name(ns);
 
-                    val_tys.push(quote! {
-                        crate::models::#ty_name
+                    ns.push_top_level(quote! {
+                        impl From<#ty_name> for #enum_name {
+                            fn from(value: #ty_name) -> Self {
+                                Self::#name(value)
+                            }
+                        }
                     });
                 }
 
+                for value in values {
+                    let ty_name = value.type_name(ns);
+
+                    if let EnumVariantTupleValue::Enum { inner, .. } = &value {
+                        ns.push_element(inner.codegen(resolved)?);
+                    }
+
+                    val_tys.push(ty_name);
+                }
+
                 Some(quote! {
                     #name(#(#val_tys),*)
                 })
@@ -144,12 +313,7 @@ impl EnumVariant {
     }
 
     pub fn codegen_display(&self) -> Option<TokenStream> {
-        let rhs = self.value.codegen_display(&self.name)?;
-        let name = format_ident!("{}", self.name);
-
-        Some(quote! {
-            Self::#name => #rhs
-        })
+        self.value.codegen_display(&self.name)
     }
 }
 
@@ -159,7 +323,6 @@ pub struct Enum {
     pub description: Option<String>,
     pub repr: Option<EnumRepr>,
     pub copy: bool,
-    pub display: bool,
     pub untagged: bool,
     pub variants: Vec<EnumVariant>,
 }
@@ -176,7 +339,6 @@ impl Enum {
         match &schema.r#enum {
             Some(OpenApiVariants::Int(int_variants)) => {
                 result.repr = Some(EnumRepr::U32);
-                result.display = true;
                 result.variants = int_variants
                     .iter()
                     .copied()
@@ -188,7 +350,6 @@ impl Enum {
                     .collect();
             }
             Some(OpenApiVariants::Str(str_variants)) => {
-                result.display = true;
                 result.variants = str_variants
                     .iter()
                     .copied()
@@ -214,7 +375,6 @@ impl Enum {
         let mut result = Self {
             name: name.to_owned(),
             copy: true,
-            display: true,
             ..Default::default()
         };
 
@@ -240,7 +400,7 @@ impl Enum {
         };
 
         for schema in schemas {
-            let value = EnumVariantTupleValue::from_schema(schema)?;
+            let value = EnumVariantTupleValue::from_schema(name, schema)?;
             let name = value.name();
 
             result.variants.push(EnumVariant {
@@ -253,7 +413,11 @@ impl Enum {
         Some(result)
     }
 
-    pub fn codegen(&self) -> Option<TokenStream> {
+    pub fn is_display(&self, resolved: &ResolvedSchema) -> bool {
+        self.variants.iter().all(|v| v.value.is_display(resolved))
+    }
+
+    pub fn codegen(&self, resolved: &ResolvedSchema) -> Option<TokenStream> {
         let repr = self.repr.map(|r| match r {
             EnumRepr::U8 => quote! { #[repr(u8)]},
             EnumRepr::U32 => quote! { #[repr(u32)]},
@@ -266,12 +430,21 @@ impl Enum {
             }
         });
 
+        let mut ns = EnumNamespace {
+            r#enum: self,
+            ident: None,
+            elements: Default::default(),
+            top_level_elements: Default::default(),
+        };
+
+        let is_display = self.is_display(resolved);
+
         let mut display = Vec::with_capacity(self.variants.len());
         let mut variants = Vec::with_capacity(self.variants.len());
         for variant in &self.variants {
-            variants.push(variant.codegen()?);
+            variants.push(variant.codegen(&mut ns, resolved)?);
 
-            if self.display {
+            if is_display {
                 display.push(variant.codegen_display()?);
             }
         }
@@ -294,7 +467,7 @@ impl Enum {
             }
         });
 
-        let display = self.display.then(|| {
+        let display = is_display.then(|| {
             quote! {
                 impl std::fmt::Display for #name {
                     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
@@ -306,6 +479,8 @@ impl Enum {
             }
         });
 
+        let module = ns.codegen();
+
         Some(quote! {
             #desc
             #[derive(Debug, Clone, PartialEq, #(#derives),*)]
@@ -315,6 +490,24 @@ impl Enum {
                 #(#variants),*
             }
             #display
+
+            #module
         })
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    use crate::openapi::schema::test::get_schema;
+
+    #[test]
+    fn is_display() {
+        let schema = get_schema();
+        let resolved = ResolvedSchema::from_open_api(&schema);
+
+        let torn_selection_name = resolved.models.get("TornSelectionName").unwrap();
+        assert!(torn_selection_name.is_display(&resolved));
+    }
+}