datagen_rs/schema/
reference.rs

1use crate::schema::transform::Transform;
2#[cfg(feature = "schema")]
3use schemars::JsonSchema;
4#[cfg(feature = "serialize")]
5use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone)]
8#[cfg_attr(feature = "schema", derive(JsonSchema))]
9#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
10#[cfg_attr(feature = "serialize", serde(rename_all = "camelCase"))]
11pub struct Reference {
12    pub reference: String,
13    pub except: Option<Vec<StringOrNumber>>,
14    pub keep_all: Option<bool>,
15    pub transform: Option<Vec<Transform>>,
16}
17
18#[derive(Debug, Clone, PartialEq)]
19#[cfg_attr(feature = "schema", derive(JsonSchema))]
20#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
21#[cfg_attr(feature = "serialize", serde(untagged, deny_unknown_fields))]
22pub enum StringOrNumber {
23    String(String),
24    Number(f64),
25}
26
27#[cfg(feature = "map-schema")]
28pub mod generate {
29    use crate::generate::current_schema::CurrentSchemaRef;
30    use crate::generate::generated_schema::generate::IntoGeneratedArc;
31    use crate::generate::generated_schema::GeneratedSchema;
32    use crate::schema::reference::{Reference, StringOrNumber};
33    use crate::schema::transform::Transform;
34    use crate::util::types::Result;
35    use rand::prelude::SliceRandom;
36    use std::sync::Arc;
37
38    impl IntoGeneratedArc for Reference {
39        fn into_generated_arc(self, schema: CurrentSchemaRef) -> Result<Arc<GeneratedSchema>> {
40            let mut reference = self.reference;
41            if !reference.starts_with("ref:") {
42                reference = format!("ref:{reference}");
43            }
44
45            let resolved = schema.resolve_ref(reference)?;
46            let Some(except) = self.except else {
47                return resolved.into_random();
48            };
49
50            let Some(resolved) = resolved.into_vec() else {
51                return Ok(Arc::new(GeneratedSchema::None));
52            };
53
54            let except = except
55                .into_iter()
56                .map(|x| match x {
57                    StringOrNumber::String(string) => {
58                        Ok(schema.resolve_ref(string)?.into_vec().unwrap_or(vec![]))
59                    }
60                    StringOrNumber::Number(number) => {
61                        Ok(vec![Arc::new(GeneratedSchema::Number(number.into()))])
62                    }
63                })
64                .collect::<Result<Vec<_>>>()?
65                .into_iter()
66                .flatten()
67                .collect::<Vec<_>>();
68
69            let resolved = resolved
70                .iter()
71                .filter(|x| !except.contains(x))
72                .cloned()
73                .collect::<Vec<_>>();
74
75            Ok(if self.keep_all.unwrap_or(false) {
76                Arc::new(GeneratedSchema::Array(resolved))
77            } else if let Some(resolved) = resolved.choose(&mut rand::thread_rng()) {
78                resolved.clone()
79            } else {
80                Arc::new(GeneratedSchema::None)
81            })
82        }
83
84        fn get_transform(&self) -> Option<Vec<Transform>> {
85            self.transform.clone()
86        }
87    }
88}