surrealml_core/storage/header/
output.rs

1//! Defines the struct housing data around the outputs of the model.
2use std::fmt;
3
4use super::normalisers::wrapper::NormaliserType;
5use crate::errors::error::{SurrealError, SurrealErrorStatus};
6use crate::safe_eject_option;
7
8/// Houses data around the outputs of the model.
9///
10/// # Fields
11/// * `name` - The name of the output.
12/// * `normaliser` - The normaliser to be applied to the output if there is one.
13#[derive(Debug, PartialEq)]
14pub struct Output {
15	pub name: Option<String>,
16	pub normaliser: Option<NormaliserType>,
17}
18
19impl Output {
20	/// Creates a new instance of the Output struct with no normaliser or name.
21	///
22	/// # Returns
23	/// A new instance of the Output struct with no normaliser or name.
24	pub fn fresh() -> Self {
25		Output {
26			name: None,
27			normaliser: None,
28		}
29	}
30
31	/// Creates a new instance of the Output struct without a normaliser.
32	///
33	/// # Arguments
34	/// * `name` - The name of the output.
35	pub fn new(name: String) -> Self {
36		Output {
37			name: Some(name),
38			normaliser: None,
39		}
40	}
41
42	/// Adds a normaliser to the output.
43	///
44	/// # Arguments
45	/// * `normaliser` - The normaliser to be applied to the output.
46	pub fn add_normaliser(&mut self, normaliser: NormaliserType) {
47		self.normaliser = Some(normaliser);
48	}
49
50	/// Converts a string to an instance of the Output struct.
51	///
52	/// # Arguments
53	/// * `data` - The string to be converted into an instance of the Output struct.
54	///
55	/// # Returns
56	/// * `Output` - The string as an instance of the Output struct.
57	pub fn from_string(data: String) -> Result<Self, SurrealError> {
58		if !data.contains("=>") {
59			return Ok(Output::fresh());
60		}
61		let mut buffer = data.split("=>");
62
63		let name = safe_eject_option!(buffer.next());
64		let name = match name {
65			"none" => None,
66			_ => Some(name.to_string()),
67		};
68
69		let normaliser = safe_eject_option!(buffer.next());
70		let normaliser = match normaliser {
71			"none" => None,
72			// Propagate malformed-normaliser errors instead of `unwrap()`-panicking on
73			// attacker-controlled header data.
74			_ => Some(NormaliserType::from_string(data)?.0),
75		};
76		Ok(Output {
77			name,
78			normaliser,
79		})
80	}
81}
82
83impl fmt::Display for Output {
84	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
85		if self.name.is_none() && self.normaliser.is_none() {
86			return write!(f, "");
87		}
88
89		let name = self.name.as_deref().unwrap_or("none");
90		let normaliser =
91			self.normaliser.as_ref().map(|n| n.to_string()).unwrap_or_else(|| "none".to_string());
92
93		write!(f, "{}=>{}", name, normaliser)
94	}
95}
96
97#[cfg(test)]
98pub mod tests {
99
100	use super::*;
101
102	#[test]
103	fn test_output_to_string() {
104		// with no normaliser
105		let mut output = Output::new("test".to_string());
106		assert_eq!(output.to_string(), "test=>none");
107
108		let normaliser_data = "a=>linear_scaling(0.0,1.0)".to_string();
109		let normaliser = NormaliserType::from_string(normaliser_data).unwrap();
110
111		output.add_normaliser(normaliser.0);
112		assert_eq!(output.to_string(), "test=>linear_scaling(0,1)");
113	}
114
115	#[test]
116	fn test_from_string() {
117		let data = "test=>linear_scaling(0,1)".to_string();
118		let output = Output::from_string(data).unwrap();
119
120		assert_eq!(output.name.unwrap(), "test");
121		assert_eq!(output.normaliser.unwrap().to_string(), "linear_scaling(0,1)");
122	}
123
124	#[test]
125	fn test_from_string_with_no_normaliser() {
126		let data = "test=>none".to_string();
127		let output = Output::from_string(data).unwrap();
128
129		assert_eq!(output.name.unwrap(), "test");
130		assert!(output.normaliser.is_none());
131	}
132
133	#[test]
134	fn test_from_string_with_no_name() {
135		let data = "none=>none".to_string();
136		let output = Output::from_string(data).unwrap();
137
138		assert!(output.name.is_none());
139		assert!(output.normaliser.is_none());
140	}
141
142	#[test]
143	fn test_from_string_with_empty_string() {
144		let data = "".to_string();
145		let output = Output::from_string(data).unwrap();
146
147		assert!(output.name.is_none());
148		assert!(output.normaliser.is_none());
149	}
150
151	#[test]
152	fn test_to_string_with_no_data() {
153		let output = Output::fresh();
154		assert_eq!(output.to_string(), "");
155	}
156
157	// Regression test for GHSA-jwr6-6444-28xv: a malformed normaliser in the output
158	// field must error rather than panic on `NormaliserType::from_string(..).unwrap()`.
159	#[test]
160	fn test_from_string_malformed_normaliser_errors() {
161		assert!(Output::from_string("col=>garbage".to_string()).is_err());
162	}
163}
surrealml_core/storage/header/output.rs

surrealml_core/storage/header/
output.rs
