Skip to main content

ad_core/plugin/
file_base.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use crate::error::ADResult;
5use crate::ndarray::NDArray;
6
7/// File write modes matching C++ NDFileMode_t.
8#[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
25/// Trait for file format writers.
26pub 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
36/// File path/name management and capture buffering for file plugins.
37pub 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    /// Construct the full file path from template/path/name/number.
71    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    /// Get the temp file path (if temp_suffix is set).
82    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    /// Create directory if needed.
92    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    /// Process an incoming array according to the current file mode.
100    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    /// Flush capture buffer: open file, write all buffered arrays, close.
136    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    /// Close stream mode.
155    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    /// Test file writer that records operations.
185    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); // auto-incremented
240    }
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        // Buffer 3 arrays
253        fb.process_array(make_array(1), &mut writer).unwrap();
254        assert_eq!(writer.writes, 0); // not flushed yet
255        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        // Should have flushed
259        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); // opened once
278        assert_eq!(writer.writes, 3);
279        assert_eq!(writer.closes, 0); // not closed yet
280
281        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        // With create_dir=0 and empty path, should be a no-op
335        fb.ensure_directory().unwrap();
336    }
337}