From 5e294b7b666366bbf4f78fce752d5fbc2a1e704f Mon Sep 17 00:00:00 2001 From: Magnus Ulimoen Date: Thu, 26 Sep 2019 11:57:28 +0200 Subject: [PATCH] support groups --- src/file.rs | 90 ++++++++++++++++++++++++++++++++++++++++++++++++---- src/group.rs | 70 ++++++++++++++++++++++++++++++++++------ tests/lib.rs | 48 +++++++++++++++++++++++----- 3 files changed, 183 insertions(+), 25 deletions(-) diff --git a/src/file.rs b/src/file.rs index 98a6fe9..f060e69 100644 --- a/src/file.rs +++ b/src/file.rs @@ -112,10 +112,11 @@ impl File { name: "root".to_string(), ncid, grpid: None, - variables: HashMap::new(), - attributes: HashMap::new(), - dimensions: HashMap::new(), - sub_groups: HashMap::new(), + variables: Default::default(), + attributes: Default::default(), + parent_dimensions: Default::default(), + dimensions: Default::default(), + groups: Default::default(), }; Ok(File { ncid, @@ -152,7 +153,7 @@ impl<'a> std::ops::DerefMut for MemFile<'a> { #[cfg(feature = "memory")] impl<'a> MemFile<'a> { pub fn new(name: Option<&str>, mem: &'a [u8]) -> error::Result { - let cstr = std::ffi::CString::new(name.unwrap_or(" ")).unwrap(); + let cstr = std::ffi::CString::new(name.unwrap_or("/")).unwrap(); let mut ncid = 0; let err; unsafe { @@ -203,6 +204,9 @@ fn get_group_dimensions(ncid: nc_type) -> error::Result error::Result error::Result error::Result> { if err != NC_NOERR { return Err(err.into()); } + if nvars == 0 { + return Ok(HashMap::new()); + } let mut varids = vec![0; nvars as usize]; let err; unsafe { @@ -406,6 +419,67 @@ fn get_variables(ncid: nc_type) -> error::Result> { Ok(variables) } +fn get_groups( + ncid: nc_type, + parent_dim: &[HashMap], +) -> error::Result> { + let err; + let mut ngroups = 0; + + unsafe { + err = nc_inq_grps(ncid, &mut ngroups, std::ptr::null_mut()); + } + if err != NC_NOERR { + return Err(err.into()); + } + if ngroups == 0 { + return Ok(HashMap::new()); + } + let err; + let mut grpids = vec![0; ngroups as usize]; + unsafe { + err = nc_inq_grps(ncid, std::ptr::null_mut(), grpids.as_mut_ptr()); + } + if err != NC_NOERR { + return Err(err.into()); + } + let mut groups = HashMap::with_capacity(ngroups as usize); + for grpid in grpids.into_iter() { + let dimensions = get_group_dimensions(grpid)?; + let variables = get_variables(grpid)?; + let mut parent_dimensions = parent_dim.to_vec(); + parent_dimensions.push(dimensions.clone()); + let subgroups = get_groups(grpid, &parent_dimensions)?; + let attributes = get_attributes(grpid, NC_GLOBAL)?; + + let mut cname = [0; NC_MAX_NAME as usize + 1]; + let err; + unsafe { + err = nc_inq_grpname(grpid, cname.as_mut_ptr()); + } + if err != NC_NOERR { + return Err(err.into()); + } + let name = unsafe { std::ffi::CStr::from_ptr(cname.as_ptr()) } + .to_string_lossy() + .to_string(); + + let g = Group { + name: name.clone(), + ncid, + grpid: Some(grpid), + attributes, + parent_dimensions: parent_dim.to_vec(), + dimensions, + variables, + groups: subgroups, + }; + groups.insert(name, g); + } + + Ok(groups) +} + fn parse_file(ncid: nc_type) -> error::Result { let _l = LOCK.lock().unwrap(); @@ -415,15 +489,17 @@ fn parse_file(ncid: nc_type) -> error::Result { let variables = get_variables(ncid)?; - let sub_groups = HashMap::new(); + // Empty parent group + let groups = get_groups(ncid, &[])?; Ok(Group { ncid, grpid: None, name: "root".into(), + parent_dimensions: Default::default(), dimensions, attributes, variables, - sub_groups, + groups, }) } diff --git a/src/group.rs b/src/group.rs index 0d8cf34..37e0e4b 100644 --- a/src/group.rs +++ b/src/group.rs @@ -13,8 +13,9 @@ pub struct Group { pub(crate) grpid: Option, pub(crate) variables: HashMap, pub(crate) attributes: HashMap, + pub(crate) parent_dimensions: Vec>, pub(crate) dimensions: HashMap, - pub(crate) sub_groups: HashMap, + pub(crate) groups: HashMap, } impl Group { @@ -36,11 +37,11 @@ impl Group { pub fn dimensions(&self) -> &HashMap { &self.dimensions } - pub fn sub_groups(&self) -> &HashMap { - &self.sub_groups + pub fn groups(&self) -> &HashMap { + &self.groups } - pub fn sub_groups_mut(&mut self, name: &str) -> Option<&mut Group> { - self.sub_groups.get_mut(name) + pub fn groups_mut(&mut self, name: &str) -> Option<&mut Group> { + self.groups.get_mut(name) } } @@ -59,10 +60,20 @@ impl Group { return Err(format!("Dimension {} already exists", name).into()); } - self.dimensions.insert( - name.into(), - Dimension::new(self.grpid.unwrap_or(self.ncid), name, len)?, - ); + let d = Dimension::new(self.grpid.unwrap_or(self.ncid), name, len)?; + self.dimensions.insert(name.into(), d.clone()); + + fn recursively_add_dim(depth: usize, name: &str, d: &Dimension, g: &mut Group) { + for (_, grp) in g.groups.iter_mut() { + grp.parent_dimensions[depth].insert(name.to_string(), d.clone()); + recursively_add_dim(depth, name, d, grp); + } + } + + let mydepth = self.parent_dimensions.len(); + for (_, grp) in self.groups.iter_mut() { + recursively_add_dim(mydepth, name, &d, grp); + } Ok(self.dimensions.get_mut(name).unwrap()) } @@ -79,7 +90,17 @@ impl Group { // Assert all dimensions exists, and get &[&Dimension] let (d, e): (Vec<_>, Vec<_>) = dims .iter() - .map(|x| self.dimensions.get(*x).ok_or(*x)) + .map(|name| { + if let Some(x) = self.dimensions.get(*name) { + return Ok(x); + } + for pdim in self.parent_dimensions.iter().rev() { + if let Some(x) = pdim.get(*name) { + return Ok(x); + } + } + Err(*name) + }) .partition(Result::is_ok); if !e.is_empty() { @@ -97,4 +118,33 @@ impl Group { Ok(self.variables.get_mut(name).unwrap()) } + + /// Add an empty group to the dataset + pub fn add_group(&mut self, name: &str) -> error::Result<&mut Group> { + let cstr = std::ffi::CString::new(name).unwrap(); + let mut grpid = 0; + let err; + unsafe { + err = nc_def_grp(self.grpid.unwrap_or(self.ncid), cstr.as_ptr(), &mut grpid); + } + if err != NC_NOERR { + return Err(err.into()); + } + + let mut parent_dimensions = self.parent_dimensions.clone(); + parent_dimensions.push(self.dimensions.clone()); + + let g = Group { + ncid: self.grpid.unwrap_or(self.ncid), + name: name.to_string(), + grpid: Some(grpid), + parent_dimensions, + attributes: Default::default(), + dimensions: Default::default(), + groups: Default::default(), + variables: Default::default(), + }; + self.groups.insert(name.to_string(), g); + Ok(self.groups.get_mut(name).unwrap()) + } } diff --git a/tests/lib.rs b/tests/lib.rs index d3602fe..9b7444a 100644 --- a/tests/lib.rs +++ b/tests/lib.rs @@ -175,23 +175,55 @@ fn open_pres_temp_4d() { } #[test] -#[cfg(feature = "ndarray")] -#[ignore] fn nc4_groups() { let f = test_location().join("simple_nc4.nc"); let file = netcdf::File::open(&f).unwrap(); - let grp1 = &file.root().sub_groups()["grp1"]; + let grp1 = &file.groups()["grp1"]; assert_eq!(grp1.name(), "grp1"); + let mut data = vec![0i32; 6 * 12]; let var = &grp1.variables().get("data").unwrap(); - let data = var.values::(None, None).unwrap(); - for x in 0..(6 * 12) { - assert_eq!(data.as_slice().unwrap()[x], x as i32); + var.values_to(&mut data, None, None).unwrap(); + for (i, x) in data.iter().enumerate() { + assert_eq!(*x, i as i32); } } +#[test] +fn create_group_dimensions() { + let d = tempfile::tempdir().unwrap(); + let filepath = d.path().join("create_group.nc"); + let mut f = netcdf::File::create(filepath).unwrap(); + + f.add_dimension("x", 20).unwrap(); + + let g = &mut f.add_group("gp1").unwrap(); + + g.add_dimension("x", 100).unwrap(); + g.add_variable::("y", &["x"]).unwrap(); + + let gg = &mut g.add_group("gp2").unwrap(); + gg.add_variable::("y", &["x"]).unwrap(); + + gg.add_dimension("x", 30).unwrap(); + gg.add_variable::("z", &["x"]).unwrap(); + + assert_eq!( + f.groups()["gp1"].variables()["y"].dimensions()[0].len(), + 100 + ); + assert_eq!( + f.groups()["gp1"].groups()["gp2"].variables()["y"].dimensions()[0].len(), + 100 + ); + assert_eq!( + f.groups()["gp1"].groups()["gp2"].variables()["z"].dimensions()[0].len(), + 30 + ); +} + // Write tests #[test] fn create() { @@ -783,8 +815,8 @@ fn read_from_memory() { (*file).root().variables()["data"] .values_to(&mut v, None, None) .unwrap(); - for i in 0..6 * 12 { - assert_eq!(v[i], i as _); + for (i, v) in v.iter().enumerate() { + assert_eq!(*v, i as _); } }