1use std::collections::BTreeMap;
4
5use super::{ProgramError, Value};
6
7#[derive(Clone, Debug, Default, PartialEq)]
27pub struct Data {
28 vars: BTreeMap<String, Value>,
29}
30
31impl Data {
32 pub fn new() -> Data {
34 Data::default()
35 }
36
37 pub fn from_json(json: &str) -> Result<Data, ProgramError> {
45 let trimmed = json.trim();
46 let mut data = Data::new();
47 if trimmed.is_empty() || trimmed == "null" {
48 return Ok(data);
49 }
50 let parsed: serde_json::Value = serde_json::from_str(trimmed)
51 .map_err(|e| ProgramError::Data(format!("data is not valid JSON: {e}")))?;
52 match parsed {
53 serde_json::Value::Array(items) => {
54 data.vars.insert("data".to_string(), array("data", &items)?);
55 }
56 serde_json::Value::Object(map) => {
57 for (name, v) in map {
58 let value = match &v {
59 serde_json::Value::Array(items) => array(&name, items)?,
60 serde_json::Value::Bool(b) => Value::Bool(*b),
61 serde_json::Value::Number(n) => {
62 if let Some(i) = n.as_i64() {
63 Value::Int(i)
64 } else if let Some(u) = n.as_u64() {
65 Value::U64(u)
66 } else {
67 Value::F64(n.as_f64().unwrap_or(f64::NAN))
68 }
69 }
70 _ => {
71 return Err(ProgramError::Data(format!(
72 "data `{name}` must be a number, a boolean, or an array of them"
73 )))
74 }
75 };
76 data.vars.insert(name, value);
77 }
78 }
79 _ => {
80 return Err(ProgramError::Data(
81 "data must be a JSON object (of arrays or scalars) or a bare array".to_string(),
82 ))
83 }
84 }
85 Ok(data)
86 }
87
88 pub fn with(mut self, name: impl Into<String>, value: impl Into<Value>) -> Data {
90 self.insert(name, value);
91 self
92 }
93
94 pub fn insert(&mut self, name: impl Into<String>, value: impl Into<Value>) -> Option<Value> {
96 self.vars.insert(name.into(), value.into())
97 }
98
99 pub fn get(&self, name: &str) -> Option<&Value> {
101 self.vars.get(name)
102 }
103
104 pub fn iter(&self) -> impl Iterator<Item = (&str, &Value)> {
106 self.vars.iter().map(|(k, v)| (k.as_str(), v))
107 }
108
109 pub fn len(&self) -> usize {
111 self.vars.len()
112 }
113
114 pub fn is_empty(&self) -> bool {
116 self.vars.is_empty()
117 }
118}
119
120fn array(name: &str, items: &[serde_json::Value]) -> Result<Value, ProgramError> {
121 items
122 .iter()
123 .map(|x| {
124 x.as_f64()
125 .or(x.as_bool().map(|b| if b { 1.0 } else { 0.0 }))
126 .ok_or_else(|| {
127 ProgramError::Data(format!("data array `{name}` must hold numbers or booleans"))
128 })
129 })
130 .collect::<Result<Vec<f64>, _>>()
131 .map(Value::from)
132}
133
134#[cfg(test)]
135mod tests {
136 use super::*;
137
138 #[test]
139 fn json_shapes() {
140 assert!(Data::from_json("").unwrap().is_empty());
141 assert!(Data::from_json(" null ").unwrap().is_empty());
142 let d = Data::from_json(r#"{"x": [1, true, 2.5], "k": 18446744073709551615, "f": 0.5}"#)
143 .unwrap();
144 assert_eq!(d.get("x"), Some(&Value::from(vec![1.0, 1.0, 2.5])));
145 assert_eq!(d.get("k"), Some(&Value::U64(u64::MAX)));
146 assert_eq!(d.get("f"), Some(&Value::F64(0.5)));
147 assert_eq!(d.len(), 3);
148 assert_eq!(
149 d.iter().map(|(k, _)| k).collect::<Vec<_>>(),
150 ["f", "k", "x"]
151 );
152 for bad in [
153 "{",
154 "3",
155 r#"{"x": "s"}"#,
156 r#"{"x": [null]}"#,
157 r#"{"x": [[1]]}"#,
158 ] {
159 assert!(
160 matches!(Data::from_json(bad), Err(ProgramError::Data(_))),
161 "{bad}"
162 );
163 }
164 }
165}