pcd_rs/
writer.rs

1//! Types for writing PCD data.
2//!
3//! [Writer](crate::writer::Writer) lets you write points sequentially to
4//! PCD file or writer given by user. The written point type must implement
5//! [PcdSerialize](crate::record::PcdSerialize) trait.
6//! See [record](crate::record) moduel doc to implement your own point type.
7#![cfg_attr(
8    feature = "derive",
9    doc = r##"
10```rust
11use eyre::Result;
12use pcd_rs::{DataKind, PcdSerialize, Writer, WriterInit};
13use std::path::Path;
14
15#[derive(PcdSerialize)]
16pub struct Point {
17    x: f32,
18    y: f32,
19    z: f32,
20}
21
22fn main() -> Result<()> {
23    let mut writer: Writer<Point, _> = WriterInit {
24        height: 300,
25        width: 1,
26        viewpoint: Default::default(),
27        data_kind: DataKind::Ascii,
28        schema: None,
29    }
30    .create("test_files/dump.pcd")?;
31
32    let point = Point {
33        x: 3.14159,
34        y: 2.71828,
35        z: -5.0,
36    };
37
38    writer.push(&point)?;
39    writer.finish()?;
40
41#  std::fs::remove_file("test_files/dump.pcd").unwrap();
42
43    Ok(())
44}
45```
46"##
47)]
48
49use crate::{
50    metas::{DataKind, FieldDef, Schema, ValueKind, ViewPoint},
51    record::{DynRecord, PcdSerialize},
52    Error, Result,
53};
54use std::{
55    collections::HashSet,
56    fs::File,
57    io::{prelude::*, BufWriter, SeekFrom},
58    marker::PhantomData,
59    path::Path,
60};
61
62/// The `DynReader` struct writes points with schema determined in runtime.
63pub type DynWriter<W> = Writer<DynRecord, W>;
64
65/// A builder type that builds [Writer](crate::writer::Writer).
66pub struct WriterInit {
67    pub width: u64,
68    pub height: u64,
69    pub viewpoint: ViewPoint,
70    pub data_kind: DataKind,
71    pub schema: Option<Schema>,
72}
73
74impl WriterInit {
75    /// Builds new [Writer](crate::writer::Writer) object from a writer.
76    /// The writer must implement both [Write](std::io::Write) and [Write](std::io::Seek)
77    /// traits.
78    pub fn build_from_writer<Record: PcdSerialize, W: Write + Seek>(
79        self,
80        writer: W,
81    ) -> Result<Writer<Record, W>, Error> {
82        let record_spec = if Record::is_dynamic() {
83            // Check if the schema is set.
84            let Some(schema) = self.schema else {
85                return Err(Error::new_invalid_writer_configuration_error(
86                    "The schema is not set on the writer. It is required for the dynamic record type."
87                ));
88            };
89
90            schema
91        } else {
92            if self.schema.is_some() {
93                return Err(Error::new_invalid_writer_configuration_error(
94                    "schema should not be set for static record type",
95                ));
96            }
97            Record::write_spec()
98        };
99
100        let seq_writer = Writer::new(
101            self.width,
102            self.height,
103            self.data_kind,
104            self.viewpoint,
105            record_spec,
106            writer,
107        )?;
108        Ok(seq_writer)
109    }
110
111    /// Builds new [Writer](crate::writer::Writer) by creating a new file.
112    pub fn create<Record, P>(self, path: P) -> Result<Writer<Record, BufWriter<File>>>
113    where
114        Record: PcdSerialize,
115        P: AsRef<Path>,
116    {
117        let writer = BufWriter::new(File::create(path.as_ref())?);
118        let seq_writer = self.build_from_writer(writer)?;
119        Ok(seq_writer)
120    }
121}
122
123/// The `Writer` struct writes points in type `T` to writer `W`.
124pub struct Writer<T, W>
125where
126    W: Write + Seek,
127{
128    data_kind: DataKind,
129    record_spec: Schema,
130    writer: W,
131    num_records: usize,
132    points_arg_begin: u64,
133    points_arg_width: usize,
134    finished: bool,
135    _phantom: PhantomData<T>,
136}
137
138impl<W, Record> Writer<Record, W>
139where
140    Record: PcdSerialize,
141    W: Write + Seek,
142{
143    fn new(
144        width: u64,
145        height: u64,
146        data_kind: DataKind,
147        viewpoint: ViewPoint,
148        record_spec: Schema,
149        mut writer: W,
150    ) -> Result<Self, Error> {
151        macro_rules! ensure {
152            ($cond:expr, $desc:expr) => {
153                if !$cond {
154                    return Err(Error::new_invalid_writer_configuration_error($desc));
155                }
156            };
157        }
158
159        // Run sanity check on the schema.
160        {
161            for FieldDef { name, count, .. } in &record_spec {
162                if name.is_empty() {}
163                ensure!(!name.is_empty(), "field name must not be empty");
164                ensure!(*count > 0, "The field count must be nonzero");
165            }
166
167            let names: HashSet<_> = record_spec.iter().map(|field| &field.name).collect();
168            ensure!(
169                names.len() == record_spec.len(),
170                "schema names must be unique"
171            );
172        }
173
174        let (points_arg_begin, points_arg_width) = {
175            let fields_args: Vec<_> = record_spec
176                .iter()
177                .map(|field| field.name.to_owned())
178                .collect();
179
180            let size_args: Vec<_> = record_spec
181                .iter()
182                .map(|field| {
183                    use ValueKind::*;
184                    let size = match field.kind {
185                        U8 | I8 => 1,
186                        U16 | I16 => 2,
187                        U32 | I32 | F32 => 4,
188                        F64 => 8,
189                    };
190                    size.to_string()
191                })
192                .collect();
193
194            let type_args: Vec<_> = record_spec
195                .iter()
196                .map(|field| {
197                    use ValueKind::*;
198                    match field.kind {
199                        U8 | U16 | U32 => "U",
200                        I8 | I16 | I32 => "I",
201                        F32 | F64 => "F",
202                    }
203                })
204                .collect();
205
206            let count_args: Vec<_> = record_spec
207                .iter()
208                .map(|field| field.count.to_string())
209                .collect();
210
211            let viewpoint_args: Vec<_> = {
212                [
213                    viewpoint.tx,
214                    viewpoint.ty,
215                    viewpoint.tz,
216                    viewpoint.qw,
217                    viewpoint.qx,
218                    viewpoint.qy,
219                    viewpoint.qz,
220                ]
221                .iter()
222                .map(|value| value.to_string())
223                .collect()
224            };
225
226            let points_arg_width = (usize::max_value() as f64).log10().floor() as usize + 1;
227
228            writeln!(writer, "# .PCD v.7 - Point Cloud Data file format")?;
229            writeln!(writer, "VERSION .7")?;
230            writeln!(writer, "FIELDS {}", fields_args.join(" "))?;
231            writeln!(writer, "SIZE {}", size_args.join(" "))?;
232            writeln!(writer, "TYPE {}", type_args.join(" "))?;
233            writeln!(writer, "COUNT {}", count_args.join(" "))?;
234            writeln!(writer, "WIDTH {}", width)?;
235            writeln!(writer, "HEIGHT {}", height)?;
236            writeln!(writer, "VIEWPOINT {}", viewpoint_args.join(" "))?;
237
238            write!(writer, "POINTS ")?;
239            let points_arg_begin = writer.seek(SeekFrom::Current(0))?;
240            writeln!(writer, "{:width$}", " ", width = points_arg_width)?;
241
242            match data_kind {
243                DataKind::Binary => writeln!(writer, "DATA binary")?,
244                DataKind::Ascii => writeln!(writer, "DATA ascii")?,
245            }
246
247            (points_arg_begin, points_arg_width)
248        };
249
250        let seq_writer = Self {
251            data_kind,
252            record_spec,
253            writer,
254            num_records: 0,
255            points_arg_begin,
256            points_arg_width,
257            finished: false,
258            _phantom: PhantomData,
259        };
260        Ok(seq_writer)
261    }
262
263    /// Finish the writer.
264    ///
265    /// The method consumes the writer must be called once when finished.
266    /// Otherwise it will panic when it drops.
267    pub fn finish(mut self) -> Result<()> {
268        self.writer.seek(SeekFrom::Start(self.points_arg_begin))?;
269        write!(
270            self.writer,
271            "{:<width$}",
272            self.num_records,
273            width = self.points_arg_width
274        )?;
275        self.finished = true;
276        Ok(())
277    }
278
279    /// Writes a new point to PCD data.
280    pub fn push(&mut self, record: &Record) -> Result<()> {
281        match self.data_kind {
282            DataKind::Binary => record.write_chunk(&mut self.writer, &self.record_spec)?,
283            DataKind::Ascii => record.write_line(&mut self.writer, &self.record_spec)?,
284        }
285
286        self.num_records += 1;
287        Ok(())
288    }
289}
290
291impl<W, Record> Drop for Writer<Record, W>
292where
293    W: Write + Seek,
294{
295    fn drop(&mut self) {
296        if !self.finished {
297            panic!("call finish() before Writer drops");
298        }
299    }
300}