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 last_written_name: String,
51}
52
53impl NDPluginFileBase {
54 pub fn new() -> Self {
55 Self {
56 file_path: String::new(),
57 file_name: String::new(),
58 file_number: 0,
59 file_template: String::new(),
60 auto_increment: false,
61 temp_suffix: String::new(),
62 create_dir: 0,
63 capture_buffer: Vec::new(),
64 num_capture: 1,
65 num_captured: 0,
66 is_open: false,
67 mode: NDFileMode::Single,
68 last_written_name: String::new(),
69 }
70 }
71
72 pub fn create_file_name(&self) -> String {
78 if self.file_template.is_empty() {
79 format!(
80 "{}{}{:04}",
81 self.file_path, self.file_name, self.file_number
82 )
83 } else {
84 let mut result = String::new();
85 let mut chars = self.file_template.chars().peekable();
86 let mut s_count = 0;
87 while let Some(c) = chars.next() {
88 if c == '%' {
89 let mut spec = String::new();
91 while let Some(&nc) = chars.peek() {
92 if nc.is_ascii_digit() || nc == '.' || nc == '-' {
93 spec.push(nc);
94 chars.next();
95 } else {
96 break;
97 }
98 }
99 match chars.next() {
100 Some('s') => {
101 s_count += 1;
102 match s_count {
103 1 => result.push_str(&self.file_path),
104 2 => result.push_str(&self.file_name),
105 _ => result.push_str(""),
106 }
107 }
108 Some('d') => {
109 let width: usize = if spec.contains('.') {
111 spec.split('.')
112 .next()
113 .and_then(|s| s.parse().ok())
114 .unwrap_or(0)
115 } else {
116 spec.parse().unwrap_or(0)
117 };
118 let precision: usize = if spec.contains('.') {
119 spec.split('.')
120 .nth(1)
121 .and_then(|s| s.parse().ok())
122 .unwrap_or(0)
123 } else {
124 0
125 };
126 let pad = width.max(precision);
127 if pad > 0 {
128 result.push_str(&format!(
129 "{:0>width$}",
130 self.file_number,
131 width = pad
132 ));
133 } else {
134 result.push_str(&self.file_number.to_string());
135 }
136 }
137 Some(other) => {
138 result.push('%');
139 result.push_str(&spec);
140 result.push(other);
141 }
142 None => result.push('%'),
143 }
144 } else {
145 result.push(c);
146 }
147 }
148 result
149 }
150 }
151
152 pub fn temp_file_path(&self) -> Option<PathBuf> {
154 if self.temp_suffix.is_empty() {
155 None
156 } else {
157 let name = self.create_file_name();
158 Some(PathBuf::from(format!("{}{}", name, self.temp_suffix)))
159 }
160 }
161
162 pub fn last_written_name(&self) -> &str {
164 &self.last_written_name
165 }
166
167 pub fn ensure_directory(&self) -> ADResult<()> {
172 if self.create_dir != 0 && !self.file_path.is_empty() {
173 std::fs::create_dir_all(&self.file_path)?;
174 }
175 Ok(())
176 }
177
178 fn write_path(&self) -> (PathBuf, Option<PathBuf>) {
180 let final_path = PathBuf::from(self.create_file_name());
181 if self.temp_suffix.is_empty() {
182 (final_path, None)
183 } else {
184 let temp = PathBuf::from(format!("{}{}", final_path.display(), self.temp_suffix));
185 (temp, Some(final_path))
186 }
187 }
188
189 fn rename_temp(temp_path: &Path, final_path: &Path) -> ADResult<()> {
191 std::fs::rename(temp_path, final_path)?;
192 Ok(())
193 }
194
195 pub fn process_array(
197 &mut self,
198 array: Arc<NDArray>,
199 writer: &mut dyn NDFileWriter,
200 ) -> ADResult<()> {
201 match self.mode {
202 NDFileMode::Single => {
203 self.last_written_name = self.create_file_name();
204 let (write_path, final_path) = self.write_path();
205 writer.open_file(&write_path, NDFileMode::Single, &array)?;
206 writer.write_file(&array)?;
207 writer.close_file()?;
208 if let Some(final_path) = final_path {
209 Self::rename_temp(&write_path, &final_path)?;
210 }
211 if self.auto_increment {
212 self.file_number += 1;
213 }
214 }
215 NDFileMode::Capture => {
216 self.capture_buffer.push(array);
217 self.num_captured = self.capture_buffer.len();
218 if self.num_captured >= self.num_capture {
219 self.flush_capture(writer)?;
220 }
221 }
222 NDFileMode::Stream => {
223 if !self.is_open {
224 self.last_written_name = self.create_file_name();
225 let (write_path, _) = self.write_path();
226 writer.open_file(&write_path, NDFileMode::Stream, &array)?;
227 self.is_open = true;
228 }
229 writer.write_file(&array)?;
230 self.num_captured += 1;
231 }
232 }
233 Ok(())
234 }
235
236 pub fn flush_capture(&mut self, writer: &mut dyn NDFileWriter) -> ADResult<()> {
238 if self.capture_buffer.is_empty() {
239 return Ok(());
240 }
241 self.last_written_name = self.create_file_name();
242 let (write_path, final_path) = self.write_path();
243 writer.open_file(&write_path, NDFileMode::Capture, &self.capture_buffer[0])?;
244 for arr in &self.capture_buffer {
245 writer.write_file(arr)?;
246 }
247 writer.close_file()?;
248 if let Some(final_path) = final_path {
249 Self::rename_temp(&write_path, &final_path)?;
250 }
251 self.capture_buffer.clear();
252 self.num_captured = 0;
253 if self.auto_increment {
254 self.file_number += 1;
255 }
256 Ok(())
257 }
258
259 pub fn close_stream(&mut self, writer: &mut dyn NDFileWriter) -> ADResult<()> {
261 if self.is_open {
262 writer.close_file()?;
263 if !self.temp_suffix.is_empty() {
265 let final_name = self.create_file_name();
266 let temp_name = format!("{}{}", final_name, self.temp_suffix);
267 Self::rename_temp(Path::new(&temp_name), Path::new(&final_name))?;
268 }
269 self.is_open = false;
270 if self.auto_increment {
271 self.file_number += 1;
272 }
273 }
274 Ok(())
275 }
276
277 pub fn is_open(&self) -> bool {
278 self.is_open
279 }
280
281 pub fn set_mode(&mut self, mode: NDFileMode) {
282 self.mode = mode;
283 }
284
285 pub fn set_num_capture(&mut self, n: usize) {
286 self.num_capture = n;
287 }
288
289 pub fn num_captured(&self) -> usize {
290 self.num_captured
291 }
292
293 pub fn mode(&self) -> NDFileMode {
294 self.mode
295 }
296
297 pub fn num_capture_target(&self) -> usize {
298 self.num_capture
299 }
300
301 pub fn capture_array(&mut self, array: Arc<NDArray>) {
302 self.capture_buffer.push(array);
303 self.num_captured = self.capture_buffer.len();
304 }
305
306 pub fn clear_capture(&mut self) {
307 self.capture_buffer.clear();
308 self.num_captured = 0;
309 }
310}
311
312#[cfg(test)]
313mod tests {
314 use super::*;
315 use crate::ndarray::{NDDataType, NDDimension};
316
317 struct MockWriter {
319 opens: Vec<PathBuf>,
320 writes: usize,
321 closes: usize,
322 multi: bool,
323 }
324
325 impl MockWriter {
326 fn new(multi: bool) -> Self {
327 Self {
328 opens: Vec::new(),
329 writes: 0,
330 closes: 0,
331 multi,
332 }
333 }
334 }
335
336 impl NDFileWriter for MockWriter {
337 fn open_file(&mut self, path: &Path, _mode: NDFileMode, _array: &NDArray) -> ADResult<()> {
338 self.opens.push(path.to_path_buf());
339 Ok(())
340 }
341 fn write_file(&mut self, _array: &NDArray) -> ADResult<()> {
342 self.writes += 1;
343 Ok(())
344 }
345 fn read_file(&mut self) -> ADResult<NDArray> {
346 Err(crate::error::ADError::UnsupportedConversion(
347 "not implemented".into(),
348 ))
349 }
350 fn close_file(&mut self) -> ADResult<()> {
351 self.closes += 1;
352 Ok(())
353 }
354 fn supports_multiple_arrays(&self) -> bool {
355 self.multi
356 }
357 }
358
359 fn make_array(id: i32) -> Arc<NDArray> {
360 let mut arr = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt8);
361 arr.unique_id = id;
362 Arc::new(arr)
363 }
364
365 #[test]
366 fn test_single_mode() {
367 let mut fb = NDPluginFileBase::new();
368 fb.file_path = "/tmp/".into();
369 fb.file_name = "test_".into();
370 fb.file_number = 1;
371 fb.auto_increment = true;
372 fb.set_mode(NDFileMode::Single);
373
374 let mut writer = MockWriter::new(false);
375 fb.process_array(make_array(1), &mut writer).unwrap();
376
377 assert_eq!(writer.opens.len(), 1);
378 assert_eq!(writer.writes, 1);
379 assert_eq!(writer.closes, 1);
380 assert_eq!(fb.file_number, 2); }
382
383 #[test]
384 fn test_capture_mode() {
385 let mut fb = NDPluginFileBase::new();
386 fb.file_path = "/tmp/".into();
387 fb.file_name = "cap_".into();
388 fb.set_mode(NDFileMode::Capture);
389 fb.set_num_capture(3);
390
391 let mut writer = MockWriter::new(true);
392
393 fb.process_array(make_array(1), &mut writer).unwrap();
395 assert_eq!(writer.writes, 0); fb.process_array(make_array(2), &mut writer).unwrap();
397 assert_eq!(writer.writes, 0);
398 fb.process_array(make_array(3), &mut writer).unwrap();
399 assert_eq!(writer.opens.len(), 1);
401 assert_eq!(writer.writes, 3);
402 assert_eq!(writer.closes, 1);
403 }
404
405 #[test]
406 fn test_stream_mode() {
407 let mut fb = NDPluginFileBase::new();
408 fb.file_path = "/tmp/".into();
409 fb.file_name = "stream_".into();
410 fb.set_mode(NDFileMode::Stream);
411
412 let mut writer = MockWriter::new(true);
413
414 fb.process_array(make_array(1), &mut writer).unwrap();
415 fb.process_array(make_array(2), &mut writer).unwrap();
416 fb.process_array(make_array(3), &mut writer).unwrap();
417
418 assert_eq!(writer.opens.len(), 1); assert_eq!(writer.writes, 3);
420 assert_eq!(writer.closes, 0); fb.close_stream(&mut writer).unwrap();
423 assert_eq!(writer.closes, 1);
424 }
425
426 #[test]
427 fn test_create_file_name_default() {
428 let mut fb = NDPluginFileBase::new();
429 fb.file_path = "/data/".into();
430 fb.file_name = "img_".into();
431 fb.file_number = 42;
432 assert_eq!(fb.create_file_name(), "/data/img_0042");
433 }
434
435 #[test]
436 fn test_create_file_name_template() {
437 let mut fb = NDPluginFileBase::new();
438 fb.file_path = "/data/".into();
439 fb.file_name = "img_".into();
440 fb.file_number = 5;
441 fb.file_template = "%s%s%d.tif".into();
442 assert_eq!(fb.create_file_name(), "/data/img_5.tif");
443 }
444
445 #[test]
446 fn test_auto_increment() {
447 let mut fb = NDPluginFileBase::new();
448 fb.file_path = "/tmp/".into();
449 fb.file_name = "t_".into();
450 fb.file_number = 0;
451 fb.auto_increment = true;
452 fb.set_mode(NDFileMode::Single);
453
454 let mut writer = MockWriter::new(false);
455 fb.process_array(make_array(1), &mut writer).unwrap();
456 assert_eq!(fb.file_number, 1);
457 fb.process_array(make_array(2), &mut writer).unwrap();
458 assert_eq!(fb.file_number, 2);
459 }
460
461 #[test]
462 fn test_temp_suffix() {
463 let mut fb = NDPluginFileBase::new();
464 fb.file_path = "/data/".into();
465 fb.file_name = "img_".into();
466 fb.file_number = 1;
467 fb.temp_suffix = ".tmp".into();
468
469 let temp = fb.temp_file_path().unwrap();
470 assert_eq!(temp.to_str().unwrap(), "/data/img_0001.tmp");
471 }
472
473 #[test]
474 fn test_ensure_directory() {
475 let fb = NDPluginFileBase::new();
476 fb.ensure_directory().unwrap();
478 }
479}