torn_api_codegen/model/
enum.rs

1use heck::ToUpperCamelCase;
2use proc_macro2::TokenStream;
3use quote::{format_ident, quote};
4
5use crate::openapi::{
6    parameter::OpenApiParameterSchema,
7    r#type::{OpenApiType, OpenApiVariants},
8};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum EnumRepr {
12    U8,
13    U32,
14}
15
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub enum EnumVariantTupleValue {
18    Ref { ty_name: String },
19    ArrayOfRefs { ty_name: String },
20}
21
22impl EnumVariantTupleValue {
23    pub fn from_schema(schema: &OpenApiType) -> Option<Self> {
24        match schema {
25            OpenApiType {
26                ref_path: Some(path),
27                ..
28            } => Some(Self::Ref {
29                ty_name: path.strip_prefix("#/components/schemas/")?.to_owned(),
30            }),
31            OpenApiType {
32                r#type: Some("array"),
33                items: Some(items),
34                ..
35            } => {
36                let OpenApiType {
37                    ref_path: Some(path),
38                    ..
39                } = items.as_ref()
40                else {
41                    return None;
42                };
43                Some(Self::ArrayOfRefs {
44                    ty_name: path.strip_prefix("#/components/schemas/")?.to_owned(),
45                })
46            }
47            _ => None,
48        }
49    }
50
51    pub fn type_name(&self) -> &str {
52        match self {
53            Self::Ref { ty_name } => ty_name,
54            Self::ArrayOfRefs { ty_name } => ty_name,
55        }
56    }
57
58    pub fn name(&self) -> String {
59        match self {
60            Self::Ref { ty_name } => ty_name.clone(),
61            Self::ArrayOfRefs { ty_name } => format!("{ty_name}s"),
62        }
63    }
64}
65
66#[derive(Debug, Clone, PartialEq, Eq)]
67pub enum EnumVariantValue {
68    Repr(u32),
69    String { rename: Option<String> },
70    Tuple(Vec<EnumVariantTupleValue>),
71}
72
73impl Default for EnumVariantValue {
74    fn default() -> Self {
75        Self::String { rename: None }
76    }
77}
78
79impl EnumVariantValue {
80    pub fn codegen_display(&self, name: &str) -> Option<TokenStream> {
81        match self {
82            Self::Repr(i) => Some(quote! { write!(f, "{}", #i) }),
83            Self::String { rename } => {
84                let name = rename.as_deref().unwrap_or(name);
85                Some(quote! { write!(f, #name) })
86            }
87            _ => None,
88        }
89    }
90}
91
92#[derive(Debug, Clone, PartialEq, Eq, Default)]
93pub struct EnumVariant {
94    pub name: String,
95    pub description: Option<String>,
96    pub value: EnumVariantValue,
97}
98
99impl EnumVariant {
100    pub fn codegen(&self) -> Option<TokenStream> {
101        let doc = self.description.as_ref().map(|d| {
102            quote! {
103                #[doc = #d]
104            }
105        });
106
107        let name = format_ident!("{}", self.name);
108
109        match &self.value {
110            EnumVariantValue::Repr(repr) => Some(quote! {
111                #doc
112                #name = #repr
113            }),
114            EnumVariantValue::String { rename } => {
115                let serde_attr = rename.as_ref().map(|r| {
116                    quote! {
117                        #[serde(rename = #r)]
118                    }
119                });
120
121                Some(quote! {
122                    #doc
123                    #serde_attr
124                    #name
125                })
126            }
127            EnumVariantValue::Tuple(values) => {
128                let mut val_tys = Vec::with_capacity(values.len());
129
130                for value in values {
131                    let ty_name = value.type_name();
132                    let ty_name = format_ident!("{ty_name}");
133
134                    val_tys.push(quote! {
135                        crate::models::#ty_name
136                    });
137                }
138
139                Some(quote! {
140                    #name(#(#val_tys),*)
141                })
142            }
143        }
144    }
145
146    pub fn codegen_display(&self) -> Option<TokenStream> {
147        let rhs = self.value.codegen_display(&self.name)?;
148        let name = format_ident!("{}", self.name);
149
150        Some(quote! {
151            Self::#name => #rhs
152        })
153    }
154}
155
156#[derive(Debug, Clone, PartialEq, Eq, Default)]
157pub struct Enum {
158    pub name: String,
159    pub description: Option<String>,
160    pub repr: Option<EnumRepr>,
161    pub copy: bool,
162    pub display: bool,
163    pub untagged: bool,
164    pub variants: Vec<EnumVariant>,
165}
166
167impl Enum {
168    pub fn from_schema(name: &str, schema: &OpenApiType) -> Option<Self> {
169        let mut result = Enum {
170            name: name.to_owned(),
171            description: schema.description.as_deref().map(ToOwned::to_owned),
172            copy: true,
173            ..Default::default()
174        };
175
176        match &schema.r#enum {
177            Some(OpenApiVariants::Int(int_variants)) => {
178                result.repr = Some(EnumRepr::U32);
179                result.display = true;
180                result.variants = int_variants
181                    .iter()
182                    .copied()
183                    .map(|i| EnumVariant {
184                        name: format!("Variant{i}"),
185                        value: EnumVariantValue::Repr(i as u32),
186                        ..Default::default()
187                    })
188                    .collect();
189            }
190            Some(OpenApiVariants::Str(str_variants)) => {
191                result.display = true;
192                result.variants = str_variants
193                    .iter()
194                    .copied()
195                    .map(|s| {
196                        let transformed = s.replace('&', "And").to_upper_camel_case();
197                        EnumVariant {
198                            value: EnumVariantValue::String {
199                                rename: (transformed != s).then(|| s.to_owned()),
200                            },
201                            name: transformed,
202                            ..Default::default()
203                        }
204                    })
205                    .collect();
206            }
207            None => return None,
208        }
209
210        Some(result)
211    }
212
213    pub fn from_parameter_schema(name: &str, schema: &OpenApiParameterSchema) -> Option<Self> {
214        let mut result = Self {
215            name: name.to_owned(),
216            copy: true,
217            display: true,
218            ..Default::default()
219        };
220
221        for var in schema.r#enum.as_ref()? {
222            let transformed = var.to_upper_camel_case();
223            result.variants.push(EnumVariant {
224                value: EnumVariantValue::String {
225                    rename: (transformed != *var).then(|| transformed.clone()),
226                },
227                name: transformed,
228                ..Default::default()
229            });
230        }
231
232        Some(result)
233    }
234
235    pub fn from_one_of(name: &str, schemas: &[OpenApiType]) -> Option<Self> {
236        let mut result = Self {
237            name: name.to_owned(),
238            untagged: true,
239            ..Default::default()
240        };
241
242        for schema in schemas {
243            let value = EnumVariantTupleValue::from_schema(schema)?;
244            let name = value.name();
245
246            result.variants.push(EnumVariant {
247                name,
248                value: EnumVariantValue::Tuple(vec![value]),
249                ..Default::default()
250            });
251        }
252
253        Some(result)
254    }
255
256    pub fn codegen(&self) -> Option<TokenStream> {
257        let repr = self.repr.map(|r| match r {
258            EnumRepr::U8 => quote! { #[repr(u8)]},
259            EnumRepr::U32 => quote! { #[repr(u32)]},
260        });
261        let name = format_ident!("{}", self.name);
262        let desc = self.description.as_ref().map(|d| {
263            quote! {
264                #repr
265                #[doc = #d]
266            }
267        });
268
269        let mut display = Vec::with_capacity(self.variants.len());
270        let mut variants = Vec::with_capacity(self.variants.len());
271        for variant in &self.variants {
272            variants.push(variant.codegen()?);
273
274            if self.display {
275                display.push(variant.codegen_display()?);
276            }
277        }
278
279        let mut derives = vec![];
280
281        if self.repr.is_some() {
282            derives.push(quote! { serde_repr::Deserialize_repr });
283        } else {
284            derives.push(quote! { serde::Deserialize });
285        }
286
287        if self.copy {
288            derives.push(quote! { Copy, Hash });
289        }
290
291        let serde_attr = self.untagged.then(|| {
292            quote! {
293                #[serde(untagged)]
294            }
295        });
296
297        let display = self.display.then(|| {
298            quote! {
299                impl std::fmt::Display for #name {
300                    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
301                        match self {
302                            #(#display),*
303                        }
304                    }
305                }
306            }
307        });
308
309        Some(quote! {
310            #desc
311            #[derive(Debug, Clone, PartialEq, #(#derives),*)]
312            #serde_attr
313            pub enum #name {
314                #(#variants),*
315            }
316            #display
317        })
318    }
319}