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 true
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 pub lazy_open: bool,
46 pub delete_driver_file: bool,
47 capture_buffer: Vec<Arc<NDArray>>,
48 num_capture: usize,
49 num_captured: usize,
50 is_open: bool,
51 mode: NDFileMode,
52 last_written_name: String,
53}
54
55impl NDPluginFileBase {
56 pub fn new() -> Self {
57 Self {
58 file_path: String::new(),
59 file_name: String::new(),
60 file_number: 0,
61 file_template: String::new(),
62 auto_increment: false,
63 temp_suffix: String::new(),
64 create_dir: 0,
65 lazy_open: false,
66 delete_driver_file: false,
67 capture_buffer: Vec::new(),
68 num_capture: 1,
69 num_captured: 0,
70 is_open: false,
71 mode: NDFileMode::Single,
72 last_written_name: String::new(),
73 }
74 }
75
76 pub fn create_file_name(&self) -> String {
82 if self.file_template.is_empty() {
83 format!(
84 "{}{}{:04}",
85 self.file_path, self.file_name, self.file_number
86 )
87 } else {
88 let mut result = String::new();
89 let mut chars = self.file_template.chars().peekable();
90 let mut s_count = 0;
91 while let Some(c) = chars.next() {
92 if c == '%' {
93 let mut spec = String::new();
95 while let Some(&nc) = chars.peek() {
96 if nc.is_ascii_digit() || nc == '.' || nc == '-' {
97 spec.push(nc);
98 chars.next();
99 } else {
100 break;
101 }
102 }
103 match chars.next() {
104 Some('s') => {
105 s_count += 1;
106 match s_count {
107 1 => result.push_str(&self.file_path),
108 2 => result.push_str(&self.file_name),
109 _ => result.push_str(""),
110 }
111 }
112 Some('d') => {
113 let width: usize = if spec.contains('.') {
115 spec.split('.')
116 .next()
117 .and_then(|s| s.parse().ok())
118 .unwrap_or(0)
119 } else {
120 spec.parse().unwrap_or(0)
121 };
122 let precision: usize = if spec.contains('.') {
123 spec.split('.')
124 .nth(1)
125 .and_then(|s| s.parse().ok())
126 .unwrap_or(0)
127 } else {
128 0
129 };
130 let pad = width.max(precision);
131 if pad > 0 {
132 result.push_str(&format!(
133 "{:0>width$}",
134 self.file_number,
135 width = pad
136 ));
137 } else {
138 result.push_str(&self.file_number.to_string());
139 }
140 }
141 Some(other) => {
142 result.push('%');
143 result.push_str(&spec);
144 result.push(other);
145 }
146 None => result.push('%'),
147 }
148 } else {
149 result.push(c);
150 }
151 }
152 result
153 }
154 }
155
156 pub fn temp_file_path(&self) -> Option<PathBuf> {
158 if self.temp_suffix.is_empty() {
159 None
160 } else {
161 let name = self.create_file_name();
162 Some(PathBuf::from(format!("{}{}", name, self.temp_suffix)))
163 }
164 }
165
166 pub fn last_written_name(&self) -> &str {
168 &self.last_written_name
169 }
170
171 pub fn ensure_directory(&self) -> ADResult<()> {
176 if self.create_dir != 0 && !self.file_path.is_empty() {
177 std::fs::create_dir_all(&self.file_path)?;
178 }
179 Ok(())
180 }
181
182 fn write_path(&self) -> (PathBuf, Option<PathBuf>) {
184 let final_path = PathBuf::from(self.create_file_name());
185 if self.temp_suffix.is_empty() {
186 (final_path, None)
187 } else {
188 let temp = PathBuf::from(format!("{}{}", final_path.display(), self.temp_suffix));
189 (temp, Some(final_path))
190 }
191 }
192
193 fn rename_temp(temp_path: &Path, final_path: &Path) -> ADResult<()> {
195 std::fs::rename(temp_path, final_path)?;
196 Ok(())
197 }
198
199 pub fn process_array(
201 &mut self,
202 array: Arc<NDArray>,
203 writer: &mut dyn NDFileWriter,
204 ) -> ADResult<()> {
205 match self.mode {
206 NDFileMode::Single => {
207 self.last_written_name = self.create_file_name();
208 let (write_path, final_path) = self.write_path();
209 writer.open_file(&write_path, NDFileMode::Single, &array)?;
210 writer.write_file(&array)?;
211 writer.close_file()?;
212 if let Some(final_path) = final_path {
213 Self::rename_temp(&write_path, &final_path)?;
214 }
215 if self.delete_driver_file {
216 if let Some(attr) = array.attributes.get("DriverFileName") {
217 let driver_file = attr.value.as_string();
218 if !driver_file.is_empty() {
219 let _ = std::fs::remove_file(&driver_file);
220 }
221 }
222 }
223 if self.auto_increment {
224 self.file_number += 1;
225 }
226 }
227 NDFileMode::Capture => {
228 self.capture_buffer.push(array);
229 self.num_captured = self.capture_buffer.len();
230 if self.num_captured >= self.num_capture {
231 self.flush_capture(writer)?;
232 }
233 }
234 NDFileMode::Stream => {
235 if !self.is_open && !self.lazy_open {
236 self.last_written_name = self.create_file_name();
237 let (write_path, _) = self.write_path();
238 writer.open_file(&write_path, NDFileMode::Stream, &array)?;
239 self.is_open = true;
240 }
241 if self.lazy_open && !self.is_open {
242 self.last_written_name = self.create_file_name();
243 let (write_path, _) = self.write_path();
244 writer.open_file(&write_path, NDFileMode::Stream, &array)?;
245 self.is_open = true;
246 }
247 writer.write_file(&array)?;
248 if self.delete_driver_file {
249 if let Some(attr) = array.attributes.get("DriverFileName") {
250 let driver_file = attr.value.as_string();
251 if !driver_file.is_empty() {
252 let _ = std::fs::remove_file(&driver_file);
253 }
254 }
255 }
256 self.num_captured += 1;
257 }
258 }
259 Ok(())
260 }
261
262 pub fn flush_capture(&mut self, writer: &mut dyn NDFileWriter) -> ADResult<()> {
269 if self.capture_buffer.is_empty() {
270 return Ok(());
271 }
272
273 if writer.supports_multiple_arrays() {
274 self.last_written_name = self.create_file_name();
276 let (write_path, final_path) = self.write_path();
277 writer.open_file(&write_path, NDFileMode::Capture, &self.capture_buffer[0])?;
278 for arr in &self.capture_buffer {
279 writer.write_file(arr)?;
280 }
281 writer.close_file()?;
282 if let Some(final_path) = final_path {
283 Self::rename_temp(&write_path, &final_path)?;
284 }
285 if self.auto_increment {
286 self.file_number += 1;
287 }
288 } else {
289 let buffer = std::mem::take(&mut self.capture_buffer);
291 for arr in &buffer {
292 self.last_written_name = self.create_file_name();
293 let (write_path, final_path) = self.write_path();
294 writer.open_file(&write_path, NDFileMode::Single, arr)?;
295 writer.write_file(arr)?;
296 writer.close_file()?;
297 if let Some(final_path) = final_path {
298 Self::rename_temp(&write_path, &final_path)?;
299 }
300 if self.auto_increment {
301 self.file_number += 1;
302 }
303 }
304 self.capture_buffer = buffer;
305 }
306
307 self.capture_buffer.clear();
308 self.num_captured = 0;
309 Ok(())
310 }
311
312 pub fn close_stream(&mut self, writer: &mut dyn NDFileWriter) -> ADResult<()> {
314 if self.is_open {
315 writer.close_file()?;
316 if !self.temp_suffix.is_empty() {
318 let final_name = self.create_file_name();
319 let temp_name = format!("{}{}", final_name, self.temp_suffix);
320 Self::rename_temp(Path::new(&temp_name), Path::new(&final_name))?;
321 }
322 self.is_open = false;
323 if self.auto_increment {
324 self.file_number += 1;
325 }
326 }
327 Ok(())
328 }
329
330 pub fn is_open(&self) -> bool {
331 self.is_open
332 }
333
334 pub fn set_mode(&mut self, mode: NDFileMode) {
335 self.mode = mode;
336 }
337
338 pub fn set_num_capture(&mut self, n: usize) {
339 self.num_capture = n;
340 }
341
342 pub fn num_captured(&self) -> usize {
343 self.num_captured
344 }
345
346 pub fn mode(&self) -> NDFileMode {
347 self.mode
348 }
349
350 pub fn num_capture_target(&self) -> usize {
351 self.num_capture
352 }
353
354 pub fn capture_array(&mut self, array: Arc<NDArray>) {
355 self.capture_buffer.push(array);
356 self.num_captured = self.capture_buffer.len();
357 }
358
359 pub fn clear_capture(&mut self) {
360 self.capture_buffer.clear();
361 self.num_captured = 0;
362 }
363}
364
365#[cfg(test)]
366mod tests {
367 use super::*;
368 use crate::ndarray::{NDDataType, NDDimension};
369
370 struct MockWriter {
372 opens: Vec<PathBuf>,
373 writes: usize,
374 closes: usize,
375 multi: bool,
376 }
377
378 impl MockWriter {
379 fn new(multi: bool) -> Self {
380 Self {
381 opens: Vec::new(),
382 writes: 0,
383 closes: 0,
384 multi,
385 }
386 }
387 }
388
389 impl NDFileWriter for MockWriter {
390 fn open_file(&mut self, path: &Path, _mode: NDFileMode, _array: &NDArray) -> ADResult<()> {
391 self.opens.push(path.to_path_buf());
392 Ok(())
393 }
394 fn write_file(&mut self, _array: &NDArray) -> ADResult<()> {
395 self.writes += 1;
396 Ok(())
397 }
398 fn read_file(&mut self) -> ADResult<NDArray> {
399 Err(crate::error::ADError::UnsupportedConversion(
400 "not implemented".into(),
401 ))
402 }
403 fn close_file(&mut self) -> ADResult<()> {
404 self.closes += 1;
405 Ok(())
406 }
407 fn supports_multiple_arrays(&self) -> bool {
408 self.multi
409 }
410 }
411
412 fn make_array(id: i32) -> Arc<NDArray> {
413 let mut arr = NDArray::new(vec![NDDimension::new(4)], NDDataType::UInt8);
414 arr.unique_id = id;
415 Arc::new(arr)
416 }
417
418 #[test]
419 fn test_single_mode() {
420 let mut fb = NDPluginFileBase::new();
421 fb.file_path = "/tmp/".into();
422 fb.file_name = "test_".into();
423 fb.file_number = 1;
424 fb.auto_increment = true;
425 fb.set_mode(NDFileMode::Single);
426
427 let mut writer = MockWriter::new(false);
428 fb.process_array(make_array(1), &mut writer).unwrap();
429
430 assert_eq!(writer.opens.len(), 1);
431 assert_eq!(writer.writes, 1);
432 assert_eq!(writer.closes, 1);
433 assert_eq!(fb.file_number, 2); }
435
436 #[test]
437 fn test_capture_mode() {
438 let mut fb = NDPluginFileBase::new();
439 fb.file_path = "/tmp/".into();
440 fb.file_name = "cap_".into();
441 fb.set_mode(NDFileMode::Capture);
442 fb.set_num_capture(3);
443
444 let mut writer = MockWriter::new(true);
445
446 fb.process_array(make_array(1), &mut writer).unwrap();
448 assert_eq!(writer.writes, 0); fb.process_array(make_array(2), &mut writer).unwrap();
450 assert_eq!(writer.writes, 0);
451 fb.process_array(make_array(3), &mut writer).unwrap();
452 assert_eq!(writer.opens.len(), 1);
454 assert_eq!(writer.writes, 3);
455 assert_eq!(writer.closes, 1);
456 }
457
458 #[test]
459 fn test_capture_mode_single_image_format() {
460 let mut fb = NDPluginFileBase::new();
461 fb.file_path = "/tmp/".into();
462 fb.file_name = "jpeg_".into();
463 fb.file_number = 0;
464 fb.auto_increment = true;
465 fb.set_mode(NDFileMode::Capture);
466 fb.set_num_capture(3);
467
468 let mut writer = MockWriter::new(false); fb.process_array(make_array(1), &mut writer).unwrap();
471 fb.process_array(make_array(2), &mut writer).unwrap();
472 fb.process_array(make_array(3), &mut writer).unwrap();
473 assert_eq!(writer.opens.len(), 3);
475 assert_eq!(writer.writes, 3);
476 assert_eq!(writer.closes, 3);
477 assert_eq!(fb.file_number, 3); }
479
480 #[test]
481 fn test_stream_mode() {
482 let mut fb = NDPluginFileBase::new();
483 fb.file_path = "/tmp/".into();
484 fb.file_name = "stream_".into();
485 fb.set_mode(NDFileMode::Stream);
486
487 let mut writer = MockWriter::new(true);
488
489 fb.process_array(make_array(1), &mut writer).unwrap();
490 fb.process_array(make_array(2), &mut writer).unwrap();
491 fb.process_array(make_array(3), &mut writer).unwrap();
492
493 assert_eq!(writer.opens.len(), 1); assert_eq!(writer.writes, 3);
495 assert_eq!(writer.closes, 0); fb.close_stream(&mut writer).unwrap();
498 assert_eq!(writer.closes, 1);
499 }
500
501 #[test]
502 fn test_create_file_name_default() {
503 let mut fb = NDPluginFileBase::new();
504 fb.file_path = "/data/".into();
505 fb.file_name = "img_".into();
506 fb.file_number = 42;
507 assert_eq!(fb.create_file_name(), "/data/img_0042");
508 }
509
510 #[test]
511 fn test_create_file_name_template() {
512 let mut fb = NDPluginFileBase::new();
513 fb.file_path = "/data/".into();
514 fb.file_name = "img_".into();
515 fb.file_number = 5;
516 fb.file_template = "%s%s%d.tif".into();
517 assert_eq!(fb.create_file_name(), "/data/img_5.tif");
518 }
519
520 #[test]
521 fn test_auto_increment() {
522 let mut fb = NDPluginFileBase::new();
523 fb.file_path = "/tmp/".into();
524 fb.file_name = "t_".into();
525 fb.file_number = 0;
526 fb.auto_increment = true;
527 fb.set_mode(NDFileMode::Single);
528
529 let mut writer = MockWriter::new(false);
530 fb.process_array(make_array(1), &mut writer).unwrap();
531 assert_eq!(fb.file_number, 1);
532 fb.process_array(make_array(2), &mut writer).unwrap();
533 assert_eq!(fb.file_number, 2);
534 }
535
536 #[test]
537 fn test_temp_suffix() {
538 let mut fb = NDPluginFileBase::new();
539 fb.file_path = "/data/".into();
540 fb.file_name = "img_".into();
541 fb.file_number = 1;
542 fb.temp_suffix = ".tmp".into();
543
544 let temp = fb.temp_file_path().unwrap();
545 assert_eq!(temp.to_str().unwrap(), "/data/img_0001.tmp");
546 }
547
548 #[test]
549 fn test_ensure_directory() {
550 let fb = NDPluginFileBase::new();
551 fb.ensure_directory().unwrap();
553 }
554}