datagen_rs/schema/
string.rs

1use crate::schema::reference::Reference;
2use crate::schema::transform::Transform;
3#[cfg(feature = "schema")]
4use schemars::JsonSchema;
5#[cfg(feature = "serialize")]
6use serde::{Deserialize, Serialize};
7use std::collections::BTreeMap;
8
9#[derive(Debug, Clone)]
10#[cfg_attr(feature = "schema", derive(JsonSchema))]
11#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
12#[cfg_attr(
13    feature = "serialize",
14    serde(deny_unknown_fields, tag = "type", rename_all = "camelCase")
15)]
16pub enum StringGenerator {
17    Uuid,
18    Email,
19    FirstName,
20    LastName,
21    FullName,
22    Username,
23    City,
24    Country,
25    CountryCode,
26    Street,
27    State,
28    ZipCode,
29    Latitude,
30    Longitude,
31    Phone,
32    #[cfg_attr(feature = "serialize", serde(rename_all = "camelCase"))]
33    Format {
34        format: String,
35        args: BTreeMap<String, FormatArg>,
36        serialize_non_strings: Option<bool>,
37    },
38}
39
40#[derive(Debug, Clone)]
41#[cfg_attr(feature = "schema", derive(JsonSchema))]
42#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
43#[cfg_attr(feature = "serialize", serde(untagged, deny_unknown_fields))]
44pub enum FormatArg {
45    String(String),
46    StringSchema(StringSchema),
47    Integer(i32),
48    Number(f64),
49    Reference(Reference),
50}
51
52#[derive(Debug, Clone)]
53#[cfg_attr(feature = "schema", derive(JsonSchema))]
54#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
55#[cfg_attr(
56    feature = "serialize",
57    serde(untagged, deny_unknown_fields, rename = "string")
58)]
59pub enum StringSchema {
60    Generated {
61        generator: StringGenerator,
62        transform: Option<Vec<Transform>>,
63    },
64    Constant {
65        value: String,
66        transform: Option<Vec<Transform>>,
67    },
68}
69
70#[cfg(feature = "generate")]
71pub mod generate {
72    use crate::generate::current_schema::CurrentSchemaRef;
73    use crate::generate::generated_schema::generate::{IntoGenerated, IntoGeneratedArc};
74    use crate::generate::generated_schema::{GeneratedSchema, IntoRandom};
75    use crate::schema::string::{FormatArg, StringGenerator, StringSchema};
76    use crate::schema::transform::Transform;
77    use crate::util::types::Result;
78    use fake::faker::address::en::{
79        CityName, CountryCode, CountryName, Latitude, Longitude, StateName, StreetName, ZipCode,
80    };
81    use fake::faker::internet::en::{FreeEmail, Username};
82    use fake::faker::name::en::{FirstName, LastName, Name};
83    use fake::faker::phone_number::en::PhoneNumber;
84    use fake::uuid::UUIDv4;
85    use fake::Fake;
86    use handlebars::Handlebars;
87    use std::collections::HashMap;
88    use std::sync::Arc;
89
90    impl IntoGeneratedArc for StringSchema {
91        fn into_generated_arc(self, schema: CurrentSchemaRef) -> Result<Arc<GeneratedSchema>> {
92            match self {
93                StringSchema::Constant { value, .. } => schema.resolve_ref(value)?.into_random(),
94                StringSchema::Generated { generator, .. } => generator.into_random(schema),
95            }
96        }
97
98        fn get_transform(&self) -> Option<Vec<Transform>> {
99            match self {
100                StringSchema::Constant { transform, .. } => transform.clone(),
101                StringSchema::Generated { transform, .. } => transform.clone(),
102            }
103        }
104    }
105
106    impl IntoGenerated for StringGenerator {
107        fn into_generated(self, schema: CurrentSchemaRef) -> Result<GeneratedSchema> {
108            Ok(match self {
109                StringGenerator::Uuid => GeneratedSchema::String(UUIDv4.fake()),
110                StringGenerator::Email => GeneratedSchema::String(FreeEmail().fake()),
111                StringGenerator::FirstName => GeneratedSchema::String(FirstName().fake()),
112                StringGenerator::LastName => GeneratedSchema::String(LastName().fake()),
113                StringGenerator::FullName => GeneratedSchema::String(Name().fake()),
114                StringGenerator::Username => GeneratedSchema::String(Username().fake()),
115                StringGenerator::City => GeneratedSchema::String(CityName().fake()),
116                StringGenerator::Country => GeneratedSchema::String(CountryName().fake()),
117                StringGenerator::CountryCode => GeneratedSchema::String(CountryCode().fake()),
118                StringGenerator::Street => GeneratedSchema::String(StreetName().fake()),
119                StringGenerator::State => GeneratedSchema::String(StateName().fake()),
120                StringGenerator::ZipCode => GeneratedSchema::String(ZipCode().fake()),
121                StringGenerator::Latitude => {
122                    GeneratedSchema::Number(Latitude().fake::<f64>().into())
123                }
124                StringGenerator::Longitude => {
125                    GeneratedSchema::Number(Longitude().fake::<f64>().into())
126                }
127                StringGenerator::Phone => GeneratedSchema::String(PhoneNumber().fake()),
128                StringGenerator::Format {
129                    format,
130                    args,
131                    serialize_non_strings,
132                } => {
133                    let mut hbs = Handlebars::new();
134                    hbs.register_template_string("template", format)?;
135
136                    let data = args
137                        .into_iter()
138                        .map(|(name, arg)| -> Result<(String, Arc<GeneratedSchema>)> {
139                            Ok((
140                                name,
141                                match arg {
142                                    FormatArg::Number(num) => {
143                                        GeneratedSchema::String(num.to_string()).into()
144                                    }
145                                    FormatArg::Integer(num) => {
146                                        GeneratedSchema::String(num.to_string()).into()
147                                    }
148                                    FormatArg::String(str) => {
149                                        schema.resolve_ref(str)?.into_random()?
150                                    }
151                                    FormatArg::StringSchema(str) => {
152                                        let res = str.into_generated_arc(schema.clone())?;
153                                        match res.as_ref() {
154                                            GeneratedSchema::Number(num) => {
155                                                GeneratedSchema::String(num.to_string()).into()
156                                            }
157                                            GeneratedSchema::Integer(num) => {
158                                                GeneratedSchema::String(num.to_string()).into()
159                                            }
160                                            _ => res,
161                                        }
162                                    }
163                                    FormatArg::Reference(reference) => {
164                                        reference.into_generated_arc(schema.clone())?
165                                    }
166                                },
167                            ))
168                        })
169                        .collect::<Result<HashMap<_, _>>>()?
170                        .into_iter()
171                        .map(|(name, arg)| -> Result<(String, String)> {
172                            if let GeneratedSchema::String(str) = arg.as_ref() {
173                                Ok((name, str.clone()))
174                            } else if serialize_non_strings
175                                .or(schema.options().serialize_non_strings)
176                                .unwrap_or(false)
177                            {
178                                Ok((name, serde_json::to_string(&arg)?))
179                            } else {
180                                Err(format!(
181                                    "Unable to format non-string value: {}",
182                                    serde_json::to_string(&arg)?
183                                )
184                                .into())
185                            }
186                        })
187                        .collect::<Result<HashMap<_, _>>>()?;
188
189                    GeneratedSchema::String(hbs.render("template", &data)?)
190                }
191            })
192        }
193
194        fn get_transform(&self) -> Option<Vec<Transform>> {
195            None
196        }
197
198        fn should_finalize(&self) -> bool {
199            false
200        }
201    }
202}