utoipa_gen/component/schema/
enum_variant.rs1use std::borrow::Cow;
2use std::marker::PhantomData;
3
4use proc_macro2::TokenStream;
5use quote::{quote, ToTokens};
6use syn::{parse_quote, TypePath};
7
8use crate::component::features::Feature;
9use crate::schema_type::SchemaType;
10use crate::Array;
11
12pub trait Variant {
13    fn to_tokens(&self) -> TokenStream;
15
16    fn get_type(&self) -> (TokenStream, TokenStream) {
18        (
19            SchemaType(&parse_quote!(str)).to_token_stream(),
20            quote! {&str},
21        )
22    }
23}
24
25pub struct SimpleEnumVariant<T: ToTokens> {
26    pub value: T,
27}
28
29impl<T> Variant for SimpleEnumVariant<T>
30where
31    T: ToTokens,
32{
33    fn to_tokens(&self) -> TokenStream {
34        self.value.to_token_stream()
35    }
36}
37
38pub struct ReprVariant<'r, T: ToTokens> {
39    pub value: T,
40    pub type_path: &'r TypePath,
41}
42
43impl<'r, T> Variant for ReprVariant<'r, T>
44where
45    T: ToTokens,
46{
47    fn to_tokens(&self) -> TokenStream {
48        self.value.to_token_stream()
49    }
50
51    fn get_type(&self) -> (TokenStream, TokenStream) {
52        (
53            SchemaType(&self.type_path.path).to_token_stream(),
54            self.type_path.to_token_stream(),
55        )
56    }
57}
58
59pub struct ObjectVariant<'o, T: ToTokens> {
60    pub item: T,
61    pub title: Option<TokenStream>,
62    pub example: Option<TokenStream>,
63    pub name: Cow<'o, str>,
64}
65
66impl<T> Variant for ObjectVariant<'_, T>
67where
68    T: ToTokens,
69{
70    fn to_tokens(&self) -> TokenStream {
71        let title = &self.title;
72        let example = &self.example;
73        let variant = &self.item;
74        let name = &self.name;
75
76        quote! {
77            utoipa::openapi::schema::ObjectBuilder::new()
78                #title
79                #example
80                .property(#name, #variant)
81                .required(#name)
82        }
83    }
84}
85
86pub struct Enum<'e, V: Variant> {
87    title: Option<TokenStream>,
88    example: Option<TokenStream>,
89    len: usize,
90    items: Array<'e, TokenStream>,
91    schema_type: TokenStream,
92    enum_type: TokenStream,
93    description: Option<TokenStream>,
94    _p: PhantomData<V>,
95}
96
97impl<V: Variant> Enum<'_, V> {
98    pub fn new<I: IntoIterator<Item = V>>(items: I) -> Self {
99        items.into_iter().collect()
100    }
101
102    pub fn with_title<I: Into<TokenStream>>(mut self, title: Option<I>) -> Self {
103        self.title = title.map(|title| title.into());
104
105        self
106    }
107
108    pub fn with_example<I: Into<TokenStream>>(mut self, example: Option<I>) -> Self {
109        self.example = example.map(|example| example.into());
110
111        self
112    }
113
114    pub fn with_description<I: Into<TokenStream>>(mut self, description: Option<I>) -> Self {
115        self.description = description.map(|description| description.into());
116
117        self
118    }
119}
120
121impl<T> ToTokens for Enum<'_, T>
122where
123    T: Variant,
124{
125    fn to_tokens(&self, tokens: &mut TokenStream) {
126        let len = &self.len;
127        let title = &self.title;
128        let example = &self.example;
129        let items = &self.items;
130        let schema_type = &self.schema_type;
131        let enum_type = &self.enum_type;
132        let description = &self.description;
133
134        tokens.extend(quote! {
135            utoipa::openapi::ObjectBuilder::new()
136                #title
137                #description
138                #example
139                .schema_type(#schema_type)
140                .enum_values::<[#enum_type; #len], #enum_type>(Some(#items))
141        })
142    }
143}
144
145impl<V: Variant> FromIterator<V> for Enum<'_, V> {
146    fn from_iter<T: IntoIterator<Item = V>>(iter: T) -> Self {
147        let mut len = 0;
148        let mut schema_type: TokenStream = quote! {};
149        let mut enum_type: TokenStream = quote! {};
150
151        let items = iter
152            .into_iter()
153            .enumerate()
154            .map(|(index, variant)| {
155                if index == 0 {
156                    (schema_type, enum_type) = variant.get_type();
157                }
158                len = index + 1;
159                variant.to_tokens()
160            })
161            .collect::<Array<TokenStream>>();
162
163        Self {
164            title: None,
165            example: None,
166            description: None,
167            len,
168            items,
169            schema_type,
170            enum_type,
171            _p: PhantomData,
172        }
173    }
174}
175
176pub struct TaggedEnum<T: Variant> {
177    items: TokenStream,
178    len: usize,
179    _p: PhantomData<T>,
180}
181
182impl<V: Variant> TaggedEnum<V> {
183    pub fn new<'t, I: IntoIterator<Item = (Cow<'t, str>, V)>>(items: I) -> Self {
184        items.into_iter().collect()
185    }
186}
187
188impl<T> ToTokens for TaggedEnum<T>
189where
190    T: Variant,
191{
192    fn to_tokens(&self, tokens: &mut TokenStream) {
193        let len = &self.len;
194        let items = &self.items;
195
196        tokens.extend(quote! {
197            Into::<utoipa::openapi::schema::OneOfBuilder>::into(utoipa::openapi::OneOf::with_capacity(#len))
198                #items
199        })
200    }
201}
202
203impl<'t, V: Variant> FromIterator<(Cow<'t, str>, V)> for TaggedEnum<V> {
204    fn from_iter<T: IntoIterator<Item = (Cow<'t, str>, V)>>(iter: T) -> Self {
205        let mut len = 0;
206
207        let items = iter
208            .into_iter()
209            .enumerate()
210            .map(|(index, (tag, variant))| {
211                len = index + 1;
212
213                let (schema_type, enum_type) = variant.get_type();
214                let item = variant.to_tokens();
215                quote! {
216                    .item(
217                        utoipa::openapi::schema::ObjectBuilder::new()
218                            .property(
219                                #tag,
220                                utoipa::openapi::schema::ObjectBuilder::new()
221                                    .schema_type(#schema_type)
222                                    .enum_values::<[#enum_type; 1], #enum_type>(Some([#item]))
223                            )
224                            .required(#tag)
225                    )
226                }
227            })
228            .collect::<TokenStream>();
229
230        Self {
231            items,
232            len,
233            _p: PhantomData,
234        }
235    }
236}
237
238pub struct UntaggedEnum {
239    title: Option<Feature>,
240}
241
242impl UntaggedEnum {
243    pub fn new() -> Self {
244        Self { title: None }
245    }
246
247    pub fn with_title(title: Option<Feature>) -> Self {
248        Self { title }
249    }
250}
251
252impl ToTokens for UntaggedEnum {
253    fn to_tokens(&self, tokens: &mut TokenStream) {
254        let title = &self.title;
255
256        tokens.extend(quote! {
257            utoipa::openapi::schema::ObjectBuilder::new()
258                .nullable(true)
259                .default(Some(serde_json::Value::Null))
260                #title
261        })
262    }
263}
264
265pub struct AdjacentlyTaggedEnum<T: Variant> {
266    items: TokenStream,
267    len: usize,
268    _p: PhantomData<T>,
269}
270
271impl<V: Variant> AdjacentlyTaggedEnum<V> {
272    pub fn new<'t, I: IntoIterator<Item = (Cow<'t, str>, Cow<'t, str>, V)>>(items: I) -> Self {
273        items.into_iter().collect()
274    }
275}
276
277impl<T> ToTokens for AdjacentlyTaggedEnum<T>
278where
279    T: Variant,
280{
281    fn to_tokens(&self, tokens: &mut TokenStream) {
282        let len = &self.len;
283        let items = &self.items;
284
285        tokens.extend(quote! {
286            Into::<utoipa::openapi::schema::OneOfBuilder>::into(utoipa::openapi::OneOf::with_capacity(#len))
287                #items
288        })
289    }
290}
291
292impl<'t, V: Variant> FromIterator<(Cow<'t, str>, Cow<'t, str>, V)> for AdjacentlyTaggedEnum<V> {
293    fn from_iter<T: IntoIterator<Item = (Cow<'t, str>, Cow<'t, str>, V)>>(iter: T) -> Self {
294        let mut len = 0;
295
296        let items = iter
297            .into_iter()
298            .enumerate()
299            .map(|(index, (tag, content, variant))| {
300                len = index + 1;
301
302                let (schema_type, enum_type) = variant.get_type();
303                let item = variant.to_tokens();
304                quote! {
305                    .item(
306                        utoipa::openapi::schema::ObjectBuilder::new()
307                            .property(
308                                #tag,
309                                utoipa::openapi::schema::ObjectBuilder::new()
310                                    .schema_type(utoipa::openapi::schema::SchemaType::String)
311                                    .enum_values::<[#enum_type; 1], #enum_type>(Some([#content]))
312                            )
313                            .required(#tag)
314                            .property(
315                                #content,
316                                utoipa::openapi::schema::ObjectBuilder::new()
317                                    .schema_type(#schema_type)
318                                    .enum_values::<[#enum_type; 1], #enum_type>(Some([#item]))
319                            )
320                            .required(#content)
321                    )
322                }
323            })
324            .collect::<TokenStream>();
325
326        Self {
327            items,
328            len,
329            _p: PhantomData,
330        }
331    }
332}
333
334pub struct CustomEnum<'c, T: ToTokens> {
338    items: T,
340    tag: Option<Cow<'c, str>>,
341}
342
343impl<'c, T: ToTokens> CustomEnum<'c, T> {
344    pub fn with_discriminator(mut self, discriminator: Option<Cow<'c, str>>) -> Self {
345        self.tag = discriminator;
346
347        self
348    }
349}
350
351impl<'c, T> ToTokens for CustomEnum<'c, T>
352where
353    T: ToTokens,
354{
355    fn to_tokens(&self, tokens: &mut TokenStream) {
356        self.items.to_tokens(tokens);
357
358        let discriminator = self.tag.as_ref().map(|tag| {
361            quote! {
362                .discriminator(Some(utoipa::openapi::schema::Discriminator::new(#tag)))
363            }
364        });
365
366        tokens.extend(quote! {
367            #discriminator
368        });
369    }
370}
371
372impl FromIterator<TokenStream> for CustomEnum<'_, TokenStream> {
373    fn from_iter<T: IntoIterator<Item = TokenStream>>(iter: T) -> Self {
374        let mut len = 0;
375
376        let items = iter
377            .into_iter()
378            .enumerate()
379            .map(|(index, variant)| {
380                len = index + 1;
381                quote! {
382                    .item(
383                        #variant
384                    )
385                }
386            })
387            .collect::<TokenStream>();
388
389        let mut tokens = TokenStream::new();
390
391        tokens.extend(quote! {
392            Into::<utoipa::openapi::schema::OneOfBuilder>::into(utoipa::openapi::OneOf::with_capacity(#len))
393                #items
394        });
395
396        CustomEnum {
397            items: tokens,
398            tag: None,
399        }
400    }
401}