tdms_rs/api/
writer.rs

1use crate::error::{Result, TdmsError};
2use crate::io::ext::TdmsWriteExt;
3use crate::model::datatypes::{DataType, PropertyValue};
4use indexmap::IndexMap;
5use std::collections::BTreeMap;
6use std::fs::File;
7use std::io::{BufWriter, Seek, Write};
8use std::marker::PhantomData;
9use std::path::{Path, PathBuf};
10
11/// A TDMS file writer.
12///
13/// Data is batched in memory and written to disk when `close()` is called.
14pub struct TdmsWriter {
15    path: PathBuf,
16    groups: IndexMap<String, WriterGroupData>,
17    properties: IndexMap<String, PropertyValue>,
18    closed: bool,
19}
20
21struct WriterGroupData {
22    name: String,
23    channels: BTreeMap<String, WriterChannelData>,
24    properties: IndexMap<String, PropertyValue>,
25}
26
27struct WriterChannelData {
28    name: String,
29    data_type: DataType,
30    data: Vec<u8>,
31    properties: IndexMap<String, PropertyValue>,
32}
33
34pub struct WriterGroup<'w> {
35    pub(crate) writer: &'w mut TdmsWriter,
36    pub(crate) group_name: String,
37}
38
39/// A typed writer handle for a single channel.
40pub struct WriterChannel<'w, T> {
41    pub(crate) writer: &'w mut TdmsWriter,
42    pub(crate) group_name: String,
43    pub(crate) channel_name: String,
44    pub(crate) _phantom: PhantomData<T>,
45}
46
47impl TdmsWriter {
48    /// Create a new TDMS writer targeting the given output path.
49    pub fn create(path: impl AsRef<Path>) -> Result<Self> {
50        Ok(Self {
51            path: path.as_ref().to_path_buf(),
52            groups: IndexMap::new(),
53            properties: IndexMap::new(),
54            closed: false,
55        })
56    }
57
58    /// Add (or look up) a group in the output file.
59    pub fn add_group(&mut self, name: impl Into<String>) -> Result<WriterGroup<'_>> {
60        let name = name.into();
61        if name.is_empty() {
62            return Err(TdmsError::InvalidName("Group name cannot be empty".into()));
63        }
64
65        if !self.groups.contains_key(&name) {
66            self.groups.insert(
67                name.clone(),
68                WriterGroupData {
69                    name: name.clone(),
70                    channels: BTreeMap::new(),
71                    properties: IndexMap::new(),
72                },
73            );
74        }
75
76        Ok(WriterGroup {
77            writer: self,
78            group_name: name,
79        })
80    }
81
82    /// Add a file-level property.
83    pub fn add_property(
84        &mut self,
85        name: impl Into<String>,
86        value: PropertyValue,
87    ) -> Result<&mut Self> {
88        let name = name.into();
89        if name.is_empty() {
90            return Err(TdmsError::InvalidName(
91                "Property name cannot be empty".into(),
92            ));
93        }
94        self.properties.insert(name, value);
95        Ok(self)
96    }
97
98    /// Write the TDMS file to disk.
99    pub fn close(mut self) -> Result<()> {
100        if self.closed {
101            return Err(TdmsError::WriterClosed);
102        }
103        self.write_file()?;
104        self.closed = true;
105        Ok(())
106    }
107
108    /// Abort writing and delete the output file if it already exists.
109    pub fn abort(self) -> Result<()> {
110        if self.path.exists() {
111            std::fs::remove_file(&self.path)?;
112        }
113        Ok(())
114    }
115
116    fn write_file(&mut self) -> Result<()> {
117        let file = File::create(&self.path)?;
118        let mut writer = BufWriter::new(file);
119
120        writer.write_all(b"TDSm")?;
121        let toc = 0x0E;
122        writer.write_u32(toc)?;
123        writer.write_u32(4712)?;
124
125        let segment_offset_pos = writer.stream_position()?;
126        writer.write_u64(0xFFFFFFFFFFFFFFFF)?;
127        writer.write_u64(0)?;
128
129        let mut object_count = 0;
130        if !self.properties.is_empty() {
131            object_count += 1;
132        }
133        for group in self.groups.values() {
134            object_count += 1;
135            object_count += group.channels.len() as u32;
136        }
137
138        writer.write_u32(object_count)?;
139
140        if !self.properties.is_empty() {
141            self.write_object_internal(&mut writer, "/", &self.properties, None)?;
142        }
143
144        for group in self.groups.values() {
145            let group_path = format!("/'{}'", group.name);
146            self.write_object_internal(&mut writer, &group_path, &group.properties, None)?;
147
148            for channel in group.channels.values() {
149                let channel_path = format!("/'{}'/'{}'", group.name, channel.name);
150                self.write_object_internal(
151                    &mut writer,
152                    &channel_path,
153                    &channel.properties,
154                    Some((&channel.data_type, channel.data.len())),
155                )?;
156            }
157        }
158
159        let raw_data_offset = writer.stream_position()?;
160        for group in self.groups.values() {
161            for channel in group.channels.values() {
162                writer.write_all(&channel.data)?;
163            }
164        }
165
166        let end_pos = writer.stream_position()?;
167        writer.seek(std::io::SeekFrom::Start(segment_offset_pos + 8))?;
168        writer.write_u64(raw_data_offset - 28)?;
169        writer.seek(std::io::SeekFrom::Start(end_pos))?;
170
171        writer.flush()?;
172        Ok(())
173    }
174
175    fn write_object_internal(
176        &self,
177        writer: &mut BufWriter<File>,
178        path: &str,
179        properties: &IndexMap<String, PropertyValue>,
180        raw_data: Option<(&DataType, usize)>,
181    ) -> Result<()> {
182        writer.write_u32(path.len() as u32)?;
183        writer.write_all(path.as_bytes())?;
184
185        let raw_data_index = if raw_data.is_some() {
186            20_u32
187        } else {
188            0xFFFFFFFF_u32
189        };
190        writer.write_u32(raw_data_index)?;
191
192        if let Some((dtype, byte_len)) = raw_data {
193            writer.write_u32(dtype.to_u32())?;
194            writer.write_u32(1)?;
195            let count = byte_len / dtype.itemsize();
196            writer.write_u64(count as u64)?;
197        }
198
199        writer.write_u32(properties.len() as u32)?;
200        for (key, value) in properties {
201            writer.write_u32(key.len() as u32)?;
202            writer.write_all(key.as_bytes())?;
203            self.write_property_value_internal(writer, value)?;
204        }
205        Ok(())
206    }
207
208    fn write_property_value_internal(
209        &self,
210        writer: &mut BufWriter<File>,
211        value: &PropertyValue,
212    ) -> Result<()> {
213        match value {
214            PropertyValue::I8(v) => {
215                writer.write_u32(DataType::I8.to_u32())?;
216                writer.write_i8(*v)?;
217            }
218            PropertyValue::I16(v) => {
219                writer.write_u32(DataType::I16.to_u32())?;
220                writer.write_i16(*v)?;
221            }
222            PropertyValue::I32(v) => {
223                writer.write_u32(DataType::I32.to_u32())?;
224                writer.write_i32(*v)?;
225            }
226            PropertyValue::I64(v) => {
227                writer.write_u32(DataType::I64.to_u32())?;
228                writer.write_i64(*v)?;
229            }
230            PropertyValue::U8(v) => {
231                writer.write_u32(DataType::U8.to_u32())?;
232                writer.write_u8(*v)?;
233            }
234            PropertyValue::U16(v) => {
235                writer.write_u32(DataType::U16.to_u32())?;
236                writer.write_u16(*v)?;
237            }
238            PropertyValue::U32(v) => {
239                writer.write_u32(DataType::U32.to_u32())?;
240                writer.write_u32(*v)?;
241            }
242            PropertyValue::U64(v) => {
243                writer.write_u32(DataType::U64.to_u32())?;
244                writer.write_u64(*v)?;
245            }
246            PropertyValue::Float(v) => {
247                writer.write_u32(DataType::Float.to_u32())?;
248                writer.write_f32(*v)?;
249            }
250            PropertyValue::Double(v) => {
251                writer.write_u32(DataType::Double.to_u32())?;
252                writer.write_f64(*v)?;
253            }
254            PropertyValue::String(s) => {
255                writer.write_u32(DataType::String.to_u32())?;
256                writer.write_u32(s.len() as u32)?;
257                writer.write_all(s.as_bytes())?;
258            }
259            PropertyValue::Boolean(b) => {
260                writer.write_u32(DataType::Boolean.to_u32())?;
261                writer.write_u8(if *b { 1 } else { 0 })?;
262            }
263            PropertyValue::TimeStamp((secs, frac)) => {
264                writer.write_u32(DataType::TimeStamp.to_u32())?;
265                writer.write_u64(*frac)?;
266                writer.write_i64(*secs)?;
267            }
268        }
269        Ok(())
270    }
271}
272
273impl<'w> WriterGroup<'w> {
274    /// Add (or look up) a channel within this group.
275    pub fn add_channel<T: WritableType>(
276        &mut self,
277        name: impl Into<String>,
278    ) -> Result<WriterChannel<'_, T>> {
279        let name = name.into();
280        if name.is_empty() {
281            return Err(TdmsError::InvalidName(
282                "Channel name cannot be empty".into(),
283            ));
284        }
285
286        let group = self
287            .writer
288            .groups
289            .get_mut(&self.group_name)
290            .ok_or_else(|| TdmsError::GroupNotFound(self.group_name.clone()))?;
291        if !group.channels.contains_key(&name) {
292            group.channels.insert(
293                name.clone(),
294                WriterChannelData {
295                    name: name.clone(),
296                    data_type: T::data_type(),
297                    data: Vec::new(),
298                    properties: IndexMap::new(),
299                },
300            );
301        }
302
303        Ok(WriterChannel {
304            writer: self.writer,
305            group_name: self.group_name.clone(),
306            channel_name: name,
307            _phantom: PhantomData,
308        })
309    }
310
311    /// Add a group-level property.
312    pub fn add_property(
313        &mut self,
314        name: impl Into<String>,
315        value: PropertyValue,
316    ) -> Result<&mut Self> {
317        let name = name.into();
318        if name.is_empty() {
319            return Err(TdmsError::InvalidName(
320                "Property name cannot be empty".into(),
321            ));
322        }
323        let group = self
324            .writer
325            .groups
326            .get_mut(&self.group_name)
327            .ok_or_else(|| TdmsError::GroupNotFound(self.group_name.clone()))?;
328        group.properties.insert(name, value);
329        Ok(self)
330    }
331}
332
333impl<'w, T: WritableType> WriterChannel<'w, T> {
334    /// Append the provided data to the channel.
335    pub fn write(&mut self, data: &[T]) -> Result<()> {
336        let group = self
337            .writer
338            .groups
339            .get_mut(&self.group_name)
340            .ok_or_else(|| TdmsError::GroupNotFound(self.group_name.clone()))?;
341        let channel = group.channels.get_mut(&self.channel_name).ok_or_else(|| {
342            TdmsError::ChannelNotFound(self.channel_name.clone(), self.group_name.clone())
343        })?;
344        T::write_to_buffer(data, &mut channel.data)?;
345        Ok(())
346    }
347
348    /// Add a channel-level property.
349    pub fn add_property(
350        &mut self,
351        name: impl Into<String>,
352        value: PropertyValue,
353    ) -> Result<&mut Self> {
354        let name = name.into();
355        if name.is_empty() {
356            return Err(TdmsError::InvalidName(
357                "Property name cannot be empty".into(),
358            ));
359        }
360        let group = self
361            .writer
362            .groups
363            .get_mut(&self.group_name)
364            .ok_or_else(|| TdmsError::GroupNotFound(self.group_name.clone()))?;
365        let channel = group.channels.get_mut(&self.channel_name).ok_or_else(|| {
366            TdmsError::ChannelNotFound(self.channel_name.clone(), self.group_name.clone())
367        })?;
368        channel.properties.insert(name, value);
369        Ok(self)
370    }
371}
372
373/// Trait implemented for element types that can be written as TDMS channel data.
374pub trait WritableType: Sized {
375    fn data_type() -> DataType;
376    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()>;
377}
378
379impl WritableType for f64 {
380    fn data_type() -> DataType {
381        DataType::Double
382    }
383    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
384        for &v in data {
385            buffer.write_f64(v)?;
386        }
387        Ok(())
388    }
389}
390
391impl WritableType for f32 {
392    fn data_type() -> DataType {
393        DataType::Float
394    }
395    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
396        for &v in data {
397            buffer.write_f32(v)?;
398        }
399        Ok(())
400    }
401}
402
403impl WritableType for i8 {
404    fn data_type() -> DataType {
405        DataType::I8
406    }
407    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
408        for &v in data {
409            buffer.write_i8(v)?;
410        }
411        Ok(())
412    }
413}
414
415impl WritableType for i16 {
416    fn data_type() -> DataType {
417        DataType::I16
418    }
419    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
420        for &v in data {
421            buffer.write_i16(v)?;
422        }
423        Ok(())
424    }
425}
426
427impl WritableType for i32 {
428    fn data_type() -> DataType {
429        DataType::I32
430    }
431    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
432        for &v in data {
433            buffer.write_i32(v)?;
434        }
435        Ok(())
436    }
437}
438
439impl WritableType for i64 {
440    fn data_type() -> DataType {
441        DataType::I64
442    }
443    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
444        for &v in data {
445            buffer.write_i64(v)?;
446        }
447        Ok(())
448    }
449}
450
451impl WritableType for u8 {
452    fn data_type() -> DataType {
453        DataType::U8
454    }
455    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
456        buffer.extend_from_slice(data);
457        Ok(())
458    }
459}
460
461impl WritableType for u16 {
462    fn data_type() -> DataType {
463        DataType::U16
464    }
465    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
466        for &v in data {
467            buffer.write_u16(v)?;
468        }
469        Ok(())
470    }
471}
472
473impl WritableType for u32 {
474    fn data_type() -> DataType {
475        DataType::U32
476    }
477    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
478        for &v in data {
479            buffer.write_u32(v)?;
480        }
481        Ok(())
482    }
483}
484
485impl WritableType for u64 {
486    fn data_type() -> DataType {
487        DataType::U64
488    }
489    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
490        for &v in data {
491            buffer.write_u64(v)?;
492        }
493        Ok(())
494    }
495}
496
497impl WritableType for bool {
498    fn data_type() -> DataType {
499        DataType::Boolean
500    }
501    fn write_to_buffer(data: &[Self], buffer: &mut Vec<u8>) -> Result<()> {
502        for &v in data {
503            buffer.write_u8(if v { 1 } else { 0 })?;
504        }
505        Ok(())
506    }
507}
tdms_rs/api/writer.rs

tdms_rs/api/
writer.rs
