conspire/io/netcdf/base/
mod.rs

1use crate::io::netcdf::{
2    NetCDF, Output, Reader, State, Writer,
3    format::{self, AttValue, Attribute, DimSpec},
4    nc_lock, reject_nul,
5};
6use std::{collections::HashMap, ffi::NulError, fs::File, io::Write};
7
8impl NetCDF {
9    pub fn close(&mut self) {
10        let _guard = nc_lock();
11        if let State::Write(writer) = &mut self.state
12            && let Some(mut output) = writer.output.take()
13        {
14            if let Some(threads) = writer.netcdf4 {
15                let mut file = std::io::BufWriter::new(&mut output.file);
16                let _ = format::hdf5::write(
17                    &writer.dims,
18                    &writer.global_attributes,
19                    &output.variables,
20                    std::mem::take(&mut output.data),
21                    threads,
22                    &mut file,
23                );
24                let _ = file.flush();
25            }
26            let _ = output.file.flush();
27        }
28    }
29    pub fn create(path: &str) -> Result<Self, NulError> {
30        Self::new(path, None)
31    }
32    /// Create a netCDF-4 (HDF5) file. `threads` bounds the pool used to compress
33    /// chunks in parallel; `0` or `1` compresses serially.
34    pub fn create_netcdf4(path: &str, threads: usize) -> Result<Self, NulError> {
35        Self::new(path, Some(threads))
36    }
37    fn new(path: &str, netcdf4: Option<usize>) -> Result<Self, NulError> {
38        reject_nul(path)?;
39        Ok(Self {
40            state: State::Write(Writer {
41                path: path.to_string(),
42                dims: Vec::new(),
43                global_attributes: Vec::new(),
44                variables: Vec::new(),
45                netcdf4,
46                output: None,
47            }),
48        })
49    }
50    pub fn open(path: &str) -> Result<Self, NulError> {
51        reject_nul(path)?;
52        let _guard = nc_lock();
53        let bytes = std::fs::read(path).expect("failed to read netCDF file");
54        let parsed = format::parse(&bytes);
55        Ok(Self {
56            state: State::Read(Reader { bytes, parsed }),
57        })
58    }
59    pub fn dimension_length(&self, name: &str) -> Result<usize, NulError> {
60        reject_nul(name)?;
61        Ok(self
62            .lookup_dimension(name)
63            .unwrap_or_else(|| panic!("no dimension named {name}")) as usize)
64    }
65    pub fn try_dimension_length(&self, name: &str) -> Result<Option<usize>, NulError> {
66        reject_nul(name)?;
67        Ok(self.lookup_dimension(name).map(|len| len as usize))
68    }
69    fn lookup_dimension(&self, name: &str) -> Option<u64> {
70        let dims: &[DimSpec] = match &self.state {
71            State::Read(reader) => &reader.parsed.dims,
72            State::Write(writer) => &writer.dims,
73        };
74        dims.iter().find(|dim| dim.name == name).map(|dim| dim.len)
75    }
76    pub fn get_variable_attribute_text(
77        &self,
78        variable: &str,
79        attr_name: &str,
80    ) -> Result<String, NulError> {
81        reject_nul(variable)?;
82        reject_nul(attr_name)?;
83        let attributes: &[Attribute] = match &self.state {
84            State::Read(reader) => reader
85                .parsed
86                .vars
87                .iter()
88                .find(|var| var.name == variable)
89                .map(|var| var.atts.as_slice()),
90            State::Write(writer) => writer
91                .variables
92                .iter()
93                .find(|var| var.name == variable)
94                .map(|var| var.attributes.as_slice()),
95        }
96        .unwrap_or_else(|| panic!("no variable named {variable}"));
97        match attributes
98            .iter()
99            .find(|attribute| attribute.name == attr_name)
100            .map(|attribute| &attribute.value)
101        {
102            Some(AttValue::Text(text)) => Ok(text.clone()),
103            _ => panic!("no text attribute {variable}::{attr_name}"),
104        }
105    }
106    pub fn define_dimension(&mut self, name: &str, len: usize) -> Result<(), NulError> {
107        reject_nul(name)?;
108        let _guard = nc_lock();
109        self.writer_defining().dims.push(DimSpec {
110            name: name.to_string(),
111            len: len as u64,
112        });
113        Ok(())
114    }
115    pub fn end_definition(&mut self) {
116        let _guard = nc_lock();
117        let writer = match &mut self.state {
118            State::Write(writer) => writer,
119            State::Read(_) => panic!("end_definition on a NetCDF opened for reading"),
120        };
121        assert!(writer.output.is_none(), "end_definition called twice");
122        let dim_index: HashMap<&str, usize> = writer
123            .dims
124            .iter()
125            .enumerate()
126            .map(|(index, dim)| (dim.name.as_str(), index))
127            .collect();
128        let mut variables: Vec<format::VarSpec> = writer
129            .variables
130            .iter_mut()
131            .map(|build| format::VarSpec {
132                name: build.name.clone(),
133                xtype: build.xtype,
134                dimids: build
135                    .dim_names
136                    .iter()
137                    .map(|dim| {
138                        *dim_index.get(dim.as_str()).unwrap_or_else(|| {
139                            panic!("variable references unknown dimension {dim}")
140                        })
141                    })
142                    .collect(),
143                atts: std::mem::take(&mut build.attributes),
144                begin: 0,
145                vsize: 0,
146                storage: format::Storage::Classic,
147            })
148            .collect();
149        let mut file = File::create(&writer.path).expect("failed to create netCDF file");
150        if writer.netcdf4.is_none() {
151            let header = format::finalize(&writer.dims, &writer.global_attributes, &mut variables);
152            file.write_all(&header)
153                .expect("failed to write netCDF header");
154        }
155        let data = vec![Vec::new(); variables.len()];
156        writer.output = Some(Output {
157            file,
158            variables,
159            data,
160        });
161    }
162    pub fn global(&mut self) {
163        let _guard = nc_lock();
164        let title = format!("{} v{}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));
165        self.writer_defining().global_attributes.extend([
166            Attribute {
167                name: "api_version".to_string(),
168                value: AttValue::Float(vec![8.25]),
169            },
170            Attribute {
171                name: "file_size".to_string(),
172                value: AttValue::Int(vec![1]),
173            },
174            Attribute {
175                name: "floating_point_word_size".to_string(),
176                value: AttValue::Int(vec![8]),
177            },
178            Attribute {
179                name: "version".to_string(),
180                value: AttValue::Float(vec![8.25]),
181            },
182            Attribute {
183                name: "title".to_string(),
184                value: AttValue::Text(title),
185            },
186        ]);
187    }
188    pub fn put_variable_attribute_text(
189        &mut self,
190        variable: &str,
191        attr_name: &str,
192        value: &str,
193    ) -> Result<(), NulError> {
194        reject_nul(variable)?;
195        reject_nul(attr_name)?;
196        reject_nul(value)?;
197        let _guard = nc_lock();
198        let build = self
199            .writer_defining()
200            .variables
201            .iter_mut()
202            .find(|build| build.name == variable)
203            .unwrap_or_else(|| panic!("no variable named {variable}"));
204        build.attributes.push(Attribute {
205            name: attr_name.to_string(),
206            value: AttValue::Text(value.to_string()),
207        });
208        Ok(())
209    }
210    pub(super) fn writer_defining(&mut self) -> &mut Writer {
211        match &mut self.state {
212            State::Write(writer) if writer.output.is_none() => writer,
213            State::Write(_) => panic!("operation not allowed after end_definition"),
214            State::Read(_) => panic!("write operation on a NetCDF opened for reading"),
215        }
216    }
217}
conspire/io/netcdf/base/mod.rs

conspire/io/netcdf/base/
mod.rs
