1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use crate::error::ADResult;
5use crate::ndarray::NDArray;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum NDFileMode {
10 Single = 0,
11 Capture = 1,
12 Stream = 2,
13}
14
15impl NDFileMode {
16 pub fn from_i32(v: i32) -> Self {
17 match v {
18 0 => Self::Single,
19 1 => Self::Capture,
20 _ => Self::Stream,
21 }
22 }
23}
24
25pub trait NDFileWriter: Send + Sync {
27 fn open_file(&mut self, path: &Path, mode: NDFileMode, array: &NDArray) -> ADResult<()>;
28 fn write_file(&mut self, array: &NDArray) -> ADResult<()>;
29 fn read_file(&mut self) -> ADResult<NDArray>;
30 fn close_file(&mut self) -> ADResult<()>;
31 fn supports_multiple_arrays(&self) -> bool {
32 false
33 }
34}
35
36pub struct NDPluginFileBase {
38 pub file_path: String,
39 pub file_name: String,
40 pub file_number: i32,
41 pub file_template: String,
42 pub auto_increment: bool,
43 pub temp_suffix: String,
44 pub create_dir: i32,
45 capture_buffer: Vec<Arc<NDArray>>,
46 num_capture: usize,
47 num_captured: usize,
48 is_open: bool,
49 mode: NDFileMode,
50}
51
52impl NDPluginFileBase {
53 pub fn new() -> Self {
54 Self {
55 file_path: String::new(),
56 file_name: String::new(),
57 file_number: 0,
58 file_template: String::new(),
59 auto_increment: true,
60 temp_suffix: String::new(),
61 create_dir: 0,
62 capture_buffer: Vec::new(),
63 num_capture: 1,
64 num_captured: 0,
65 is_open: false,
66 mode: NDFileMode::Single,
67 }
68 }
69
70 pub fn create_file_name(&self) -> String {
72 if self.file_template.is_empty() {
73 format!("{}{}{:04}", self.file_path, self.file_name, self.file_number)
74 } else {
75 self.file_template
76 .replace("%s%s", &format!("{}{}", self.file_path, self.file_name))
77 .replace("%d", &self.file_number.to_string())
78 }
79 }
80
81 pub fn temp_file_path(&self) -> Option<PathBuf> {
83 if self.temp_suffix.is_empty() {
84 None
85 } else {
86 let name = self.create_file_name();
87 Some(PathBuf::from(format!("{}{}", name, self.temp_suffix)))
88 }
89 }
90
91 pub fn ensure_directory(&self) -> ADResult<()> {
93 if self.create_dir > 0 && !self.file_path.is_empty() {
94 std::fs::create_dir_all(&self.file_path)?;
95 }
96 Ok(())
97 }
98
99 pub fn process_array(
101 &mut self,
102 array: Arc<NDArray>,
103 writer: &mut dyn NDFileWriter,
104 ) -> ADResult<()> {
105 match self.mode {
106 NDFileMode::Single => {
107 let path = PathBuf::from(self.create_file_name());
108 writer.open_file(&path, NDFileMode::Single, &array)?;
109 writer.write_file(&array)?;
110 writer.close_file()?;
111 if self.auto_increment {
112 self.file_number += 1;
113 }
114 }
115 NDFileMode::Capture => {
116 self.capture_buffer.push(array);
117 self.num_captured = self.capture_buffer.len();
118 if self.num_captured >= self.num_capture {
119 self.flush_capture(writer)?;
120 }
121 }
122 NDFileMode::Stream => {
123 if !self.is_open {
124 let path = PathBuf::from(self.create_file_name());
125 writer.open_file(&path, NDFileMode::Stream, &array)?;
126 self.is_open = true;
127 }
128 writer.write_file(&array)?;
129 self.num_captured += 1;
130 }
131 }
132 Ok(())
133 }
134
135 pub fn flush_capture(&mut self, writer: &mut dyn NDFileWriter) -> ADResult<()> {
137 if self.capture_buffer.is_empty() {
138 return Ok(());
139 }
140 let path = PathBuf::from(self.create_file_name());
141 writer.open_file(&path, NDFileMode::Capture, &self.capture_buffer[0])?;
142 for arr in &self.capture_buffer {
143 writer.write_file(arr)?;
144 }
145 writer.close_file()?;
146 self.capture_buffer.clear();
147 self.num_captured = 0;
148 if self.auto_increment {
149 self.file_number += 1;
150 }
151 Ok(())
152 }
153
154 pub fn close_stream(&mut self, writer: &mut dyn NDFileWriter) -> ADResult<()> {
156 if self.is_open {
157 writer.close_file()?;
158 self.is_open = false;
159 if self.auto_increment {
160 self.file_number += 1;
161 }
162 }
163 Ok(())
164 }
165
166 pub fn set_mode(&mut self, mode: NDFileMode) {
167 self.mode = mode;
168 }
169
170 pub fn set_num_capture(&mut self, n: usize) {
171 self.num_capture = n;
172 }
173
174 pub fn num_captured(&self) -> usize {
175 self.num_captured
176 }
177}
178
179#[cfg(test)]
180mod tests {
181 use super::*;
182 use crate::ndarray::{NDDataType, NDDimension};
183
184 struct MockWriter {
186 opens: Vec<PathBuf>,
187 writes: usize,
188 closes: usize,
189 multi: bool,
190 }
191
192 impl MockWriter {
193 fn new(multi: bool) -> Self {
194 Self { opens: Vec::new(), writes: 0, closes: 0, multi }
195 }
196 }
197
198 impl NDFileWriter for MockWriter {
199 fn open_file(&mut self, path: &Path, _mode: NDFileMode, _array: &NDArray) -> ADResult<()> {
200 self.opens.push(path.to_path_buf());
201 Ok(())
202 }
203 fn write_file(&mut self, _array: &NDArray) -> ADResult<()> {
204 self.writes += 1;
205 Ok(())
206 }
207 fn read_file(&mut self) -> ADResult<NDArray> {
208 Err(crate::error::ADError::UnsupportedConversion("not implemented".into()))
209 }
210 fn close_file(&mut self) -> ADResult<()> {
211 self.closes += 1;
212 Ok(())
213 }
214 fn supports_multiple_arrays(&self) -> bool {
215 self.multi
216 }
217 }
218
219 fn make_array(id: i32) -> Arc<NDArray> {
220 let mut arr = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt8);
221 arr.unique_id = id;
222 Arc::new(arr)
223 }
224
225 #[test]
226 fn test_single_mode() {
227 let mut fb = NDPluginFileBase::new();
228 fb.file_path = "/tmp/".into();
229 fb.file_name = "test_".into();
230 fb.file_number = 1;
231 fb.set_mode(NDFileMode::Single);
232
233 let mut writer = MockWriter::new(false);
234 fb.process_array(make_array(1), &mut writer).unwrap();
235
236 assert_eq!(writer.opens.len(), 1);
237 assert_eq!(writer.writes, 1);
238 assert_eq!(writer.closes, 1);
239 assert_eq!(fb.file_number, 2); }
241
242 #[test]
243 fn test_capture_mode() {
244 let mut fb = NDPluginFileBase::new();
245 fb.file_path = "/tmp/".into();
246 fb.file_name = "cap_".into();
247 fb.set_mode(NDFileMode::Capture);
248 fb.set_num_capture(3);
249
250 let mut writer = MockWriter::new(true);
251
252 fb.process_array(make_array(1), &mut writer).unwrap();
254 assert_eq!(writer.writes, 0); fb.process_array(make_array(2), &mut writer).unwrap();
256 assert_eq!(writer.writes, 0);
257 fb.process_array(make_array(3), &mut writer).unwrap();
258 assert_eq!(writer.opens.len(), 1);
260 assert_eq!(writer.writes, 3);
261 assert_eq!(writer.closes, 1);
262 }
263
264 #[test]
265 fn test_stream_mode() {
266 let mut fb = NDPluginFileBase::new();
267 fb.file_path = "/tmp/".into();
268 fb.file_name = "stream_".into();
269 fb.set_mode(NDFileMode::Stream);
270
271 let mut writer = MockWriter::new(true);
272
273 fb.process_array(make_array(1), &mut writer).unwrap();
274 fb.process_array(make_array(2), &mut writer).unwrap();
275 fb.process_array(make_array(3), &mut writer).unwrap();
276
277 assert_eq!(writer.opens.len(), 1); assert_eq!(writer.writes, 3);
279 assert_eq!(writer.closes, 0); fb.close_stream(&mut writer).unwrap();
282 assert_eq!(writer.closes, 1);
283 }
284
285 #[test]
286 fn test_create_file_name_default() {
287 let mut fb = NDPluginFileBase::new();
288 fb.file_path = "/data/".into();
289 fb.file_name = "img_".into();
290 fb.file_number = 42;
291 assert_eq!(fb.create_file_name(), "/data/img_0042");
292 }
293
294 #[test]
295 fn test_create_file_name_template() {
296 let mut fb = NDPluginFileBase::new();
297 fb.file_path = "/data/".into();
298 fb.file_name = "img_".into();
299 fb.file_number = 5;
300 fb.file_template = "%s%s%d.tif".into();
301 assert_eq!(fb.create_file_name(), "/data/img_5.tif");
302 }
303
304 #[test]
305 fn test_auto_increment() {
306 let mut fb = NDPluginFileBase::new();
307 fb.file_path = "/tmp/".into();
308 fb.file_name = "t_".into();
309 fb.file_number = 0;
310 fb.set_mode(NDFileMode::Single);
311
312 let mut writer = MockWriter::new(false);
313 fb.process_array(make_array(1), &mut writer).unwrap();
314 assert_eq!(fb.file_number, 1);
315 fb.process_array(make_array(2), &mut writer).unwrap();
316 assert_eq!(fb.file_number, 2);
317 }
318
319 #[test]
320 fn test_temp_suffix() {
321 let mut fb = NDPluginFileBase::new();
322 fb.file_path = "/data/".into();
323 fb.file_name = "img_".into();
324 fb.file_number = 1;
325 fb.temp_suffix = ".tmp".into();
326
327 let temp = fb.temp_file_path().unwrap();
328 assert_eq!(temp.to_str().unwrap(), "/data/img_0001.tmp");
329 }
330
331 #[test]
332 fn test_ensure_directory() {
333 let fb = NDPluginFileBase::new();
334 fb.ensure_directory().unwrap();
336 }
337}