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}