From 765d55ce365f66f94419cf08e0503101b3abe22b Mon Sep 17 00:00:00 2001 From: Jakob Kraus <52459467+JabobKrauskopf@users.noreply.github.com> Date: Sat, 13 Dec 2025 05:17:24 +0100 Subject: [PATCH] fix: attributes_mut validaing against schema --- .../src/medrecord/attributes.rs | 180 ++++++++++++++++++ crates/medmodels-core/src/medrecord/mod.rs | 45 +++-- rustmodels/src/medrecord/mod.rs | 86 ++++----- 3 files changed, 245 insertions(+), 66 deletions(-) create mode 100644 crates/medmodels-core/src/medrecord/attributes.rs diff --git a/crates/medmodels-core/src/medrecord/attributes.rs b/crates/medmodels-core/src/medrecord/attributes.rs new file mode 100644 index 00000000..f36d2b4b --- /dev/null +++ b/crates/medmodels-core/src/medrecord/attributes.rs @@ -0,0 +1,180 @@ +use crate::{ + errors::MedRecordError, + prelude::{ + Attributes, EdgeIndex, Group, MedRecordAttribute, MedRecordValue, NodeIndex, SchemaType, + }, + MedRecord, +}; + +macro_rules! impl_attributes_mut { + ( + $struct_name:ident, + $index_type:ty, + $index_field:ident, + $entity:literal, + $contains_fn:ident, + $groups_of_fn:ident, + $get_attributes_fn:ident, + $get_attributes_mut_fn:ident, + $schema_update_fn:ident, + $schema_validate_fn:ident + ) => { + pub struct $struct_name<'a> { + $index_field: &'a $index_type, + medrecord: &'a mut MedRecord, + } + + impl<'a> $struct_name<'a> { + pub(crate) fn new( + $index_field: &'a $index_type, + medrecord: &'a mut MedRecord, + ) -> Result { + if !medrecord.$contains_fn($index_field) { + return Err(MedRecordError::IndexError(format!( + concat!("Cannot find ", $entity, " with index {}"), + $index_field + ))); + } + + Ok(Self { + $index_field, + medrecord, + }) + } + + fn get_groups(&self) -> Vec { + self.medrecord + .$groups_of_fn(self.$index_field) + .expect(concat!($entity, " must exist.")) + .cloned() + .collect() + } + + fn handle_schema( + &mut self, + attributes: &Attributes, + groups: &[Group], + ) -> Result<(), MedRecordError> { + let schema = &mut self.medrecord.schema; + + match schema.schema_type() { + SchemaType::Inferred => { + if groups.is_empty() { + schema.$schema_update_fn(attributes, None, false); + } else { + for group in groups { + schema.$schema_update_fn(attributes, Some(group), false); + } + } + } + SchemaType::Provided => { + if groups.is_empty() { + schema.$schema_validate_fn(self.$index_field, attributes, None)?; + } else { + for group in groups { + schema.$schema_validate_fn( + self.$index_field, + attributes, + Some(group), + )?; + } + } + } + } + + Ok(()) + } + + fn set_attributes(&mut self, attributes: Attributes) { + *self + .medrecord + .graph + .$get_attributes_mut_fn(self.$index_field) + .expect(concat!($entity, " must exist.")) = attributes; + } + + pub fn replace_attributes( + &mut self, + attributes: Attributes, + ) -> Result<(), MedRecordError> { + let groups = self.get_groups(); + self.handle_schema(&attributes, &groups)?; + self.set_attributes(attributes); + Ok(()) + } + + pub fn update_attribute( + &mut self, + attribute: &MedRecordAttribute, + value: MedRecordValue, + ) -> Result<(), MedRecordError> { + let groups = self.get_groups(); + + let mut attributes = self + .medrecord + .$get_attributes_fn(self.$index_field) + .expect(concat!($entity, " must exist.")) + .clone(); + attributes + .entry(attribute.clone()) + .and_modify(|v| *v = value.clone()) + .or_insert(value); + + self.handle_schema(&attributes, &groups)?; + self.set_attributes(attributes); + Ok(()) + } + + pub fn remove_attribute( + &mut self, + attribute: &MedRecordAttribute, + ) -> Result { + let groups = self.get_groups(); + + let mut attributes = self + .medrecord + .$get_attributes_fn(self.$index_field) + .expect(concat!($entity, " must exist.")) + .clone(); + let removed_value = attributes.remove(attribute); + + let Some(removed_value) = removed_value else { + return Err(MedRecordError::KeyError(format!( + concat!("Attribute {} does not exist on ", $entity, " {}"), + attribute, self.$index_field + ))); + }; + + self.handle_schema(&attributes, &groups)?; + self.set_attributes(attributes); + Ok(removed_value) + } + } + }; +} + +impl_attributes_mut!( + NodeAttributesMut, + NodeIndex, + node_index, + "node", + contains_node, + groups_of_node, + node_attributes, + node_attributes_mut, + update_node, + validate_node +); + +impl_attributes_mut!( + EdgeAttributesMut, + EdgeIndex, + edge_index, + "edge", + contains_edge, + groups_of_edge, + edge_attributes, + edge_attributes_mut, + update_edge, + validate_edge +); diff --git a/crates/medmodels-core/src/medrecord/mod.rs b/crates/medmodels-core/src/medrecord/mod.rs index dcfedfef..4dc4bab3 100644 --- a/crates/medmodels-core/src/medrecord/mod.rs +++ b/crates/medmodels-core/src/medrecord/mod.rs @@ -1,3 +1,4 @@ +pub mod attributes; pub mod datatypes; mod example_dataset; mod graph; @@ -15,6 +16,7 @@ pub use self::{ use crate::{ errors::MedRecordError, medrecord::{ + attributes::{EdgeAttributesMut, NodeAttributesMut}, overview::{GroupOverview, Overview, DEFAULT_TRUNCATE_DETAILS}, polars::DataFramesExport, }, @@ -402,13 +404,11 @@ impl MedRecord { .map_err(MedRecordError::from) } - pub fn node_attributes_mut( - &mut self, - node_index: &NodeIndex, - ) -> Result<&mut Attributes, MedRecordError> { - self.graph - .node_attributes_mut(node_index) - .map_err(MedRecordError::from) + pub fn node_attributes_mut<'a>( + &'a mut self, + node_index: &'a NodeIndex, + ) -> Result, MedRecordError> { + NodeAttributesMut::new(node_index, self) } pub fn outgoing_edges( @@ -439,13 +439,11 @@ impl MedRecord { .map_err(MedRecordError::from) } - pub fn edge_attributes_mut( - &mut self, - edge_index: &EdgeIndex, - ) -> Result<&mut Attributes, MedRecordError> { - self.graph - .edge_attributes_mut(edge_index) - .map_err(MedRecordError::from) + pub fn edge_attributes_mut<'a>( + &'a mut self, + edge_index: &'a EdgeIndex, + ) -> Result, MedRecordError> { + EdgeAttributesMut::new(edge_index, self) } pub fn edge_endpoints( @@ -1314,17 +1312,18 @@ mod test { fn test_node_attributes_mut() { let mut medrecord = create_medrecord(); - let attributes = medrecord.node_attributes_mut(&"0".into()).unwrap(); - - assert_eq!(&create_nodes()[0].1, attributes); + let node_index = "0".into(); + let mut attributes = medrecord.node_attributes_mut(&node_index).unwrap(); let new_attributes = HashMap::from([("0".into(), "1".into()), ("2".into(), "3".into())]); - attributes.clone_from(&new_attributes); + attributes + .replace_attributes(new_attributes.clone()) + .unwrap(); assert_eq!( &new_attributes, - medrecord.node_attributes(&"0".into()).unwrap() + medrecord.node_attributes(&node_index).unwrap() ); } @@ -1407,13 +1406,13 @@ mod test { fn test_edge_attributes_mut() { let mut medrecord = create_medrecord(); - let attributes = medrecord.edge_attributes_mut(&0).unwrap(); - - assert_eq!(&create_edges()[0].2, attributes); + let mut attributes = medrecord.edge_attributes_mut(&0).unwrap(); let new_attributes = HashMap::from([("0".into(), "1".into()), ("2".into(), "3".into())]); - attributes.clone_from(&new_attributes); + attributes + .replace_attributes(new_attributes.clone()) + .unwrap(); assert_eq!(&new_attributes, medrecord.edge_attributes(&0).unwrap()); } diff --git a/rustmodels/src/medrecord/mod.rs b/rustmodels/src/medrecord/mod.rs index d7f39ab9..f89379c3 100644 --- a/rustmodels/src/medrecord/mod.rs +++ b/rustmodels/src/medrecord/mod.rs @@ -15,9 +15,8 @@ use crate::{ }; use attribute::PyMedRecordAttribute; use errors::PyMedRecordError; -use medmodels::core::{ - errors::MedRecordError, - medrecord::{Attributes, EdgeIndex, MedRecord, MedRecordAttribute, MedRecordValue}, +use medmodels::core::medrecord::{ + Attributes, EdgeIndex, MedRecord, MedRecordAttribute, MedRecordValue, }; use pyo3::{ prelude::*, @@ -339,9 +338,9 @@ impl PyMedRecord { pub fn remove_nodes( &mut self, - node_index: Vec, + node_indices: Vec, ) -> PyResult> { - node_index + node_indices .into_iter() .map(|node_index| { let attributes = self @@ -356,18 +355,20 @@ impl PyMedRecord { pub fn replace_node_attributes( &mut self, - node_index: Vec, + node_indices: Vec, attributes: PyAttributes, ) -> PyResult<()> { let attributes: Attributes = attributes.deep_into(); - for node_index in node_index { - let current_attributes = self + for node_index in node_indices { + let mut current_attributes = self .0 .node_attributes_mut(&node_index) .map_err(PyMedRecordError::from)?; - current_attributes.clone_from(&attributes); + current_attributes + .replace_attributes(attributes.clone()) + .map_err(PyMedRecordError::from)?; } Ok(()) @@ -375,23 +376,22 @@ impl PyMedRecord { pub fn update_node_attribute( &mut self, - node_index: Vec, + node_indices: Vec, attribute: PyMedRecordAttribute, value: PyMedRecordValue, ) -> PyResult<()> { let attribute: MedRecordAttribute = attribute.into(); let value: MedRecordValue = value.into(); - for node_index in node_index { - let node_attributes = self + for node_index in node_indices { + let mut node_attributes = self .0 .node_attributes_mut(&node_index) .map_err(PyMedRecordError::from)?; node_attributes - .entry(attribute.clone()) - .or_default() - .clone_from(&value) + .update_attribute(&attribute, value.clone()) + .map_err(PyMedRecordError::from)?; } Ok(()) @@ -399,23 +399,20 @@ impl PyMedRecord { pub fn remove_node_attribute( &mut self, - node_index: Vec, + node_indices: Vec, attribute: PyMedRecordAttribute, ) -> PyResult<()> { let attribute: MedRecordAttribute = attribute.into(); - for node_index in node_index { - let node_attributes = self + for node_index in node_indices { + let mut node_attributes = self .0 .node_attributes_mut(&node_index) .map_err(PyMedRecordError::from)?; node_attributes - .remove(&attribute) - .ok_or(PyMedRecordError::from(MedRecordError::KeyError(format!( - "Cannot find attribute {} in node {}", - attribute, *node_index - ))))?; + .remove_attribute(&attribute) + .map_err(PyMedRecordError::from)?; } Ok(()) @@ -440,9 +437,9 @@ impl PyMedRecord { pub fn remove_edges( &mut self, - edge_index: Vec, + edge_indices: Vec, ) -> PyResult> { - edge_index + edge_indices .into_iter() .map(|edge_index| { let attributes = self @@ -457,18 +454,20 @@ impl PyMedRecord { pub fn replace_edge_attributes( &mut self, - edge_index: Vec, + edge_indices: Vec, attributes: PyAttributes, ) -> PyResult<()> { let attributes: Attributes = attributes.deep_into(); - for edge_index in edge_index { - let current_attributes = self + for edge_index in edge_indices { + let mut current_attributes = self .0 .edge_attributes_mut(&edge_index) .map_err(PyMedRecordError::from)?; - current_attributes.clone_from(&attributes); + current_attributes + .replace_attributes(attributes.clone()) + .map_err(PyMedRecordError::from)?; } Ok(()) @@ -476,20 +475,22 @@ impl PyMedRecord { pub fn update_edge_attribute( &mut self, - edge_index: Vec, + edge_indices: Vec, attribute: PyMedRecordAttribute, value: PyMedRecordValue, ) -> PyResult<()> { - for edge_index in edge_index { - let edge_attributes = self + let attribute: MedRecordAttribute = attribute.into(); + let value: MedRecordValue = value.into(); + + for edge_index in edge_indices { + let mut edge_attributes = self .0 .edge_attributes_mut(&edge_index) .map_err(PyMedRecordError::from)?; edge_attributes - .entry((*attribute).clone()) - .or_default() - .clone_from(&value); + .update_attribute(&attribute, value.clone()) + .map_err(PyMedRecordError::from)?; } Ok(()) @@ -497,21 +498,20 @@ impl PyMedRecord { pub fn remove_edge_attribute( &mut self, - edge_index: Vec, + edge_indices: Vec, attribute: PyMedRecordAttribute, ) -> PyResult<()> { - for edge_index in edge_index { - let edge_attributes = self + let attribute: MedRecordAttribute = attribute.into(); + + for edge_index in edge_indices { + let mut edge_attributes = self .0 .edge_attributes_mut(&edge_index) .map_err(PyMedRecordError::from)?; edge_attributes - .remove(&attribute) - .ok_or(PyMedRecordError::from(MedRecordError::KeyError(format!( - "Cannot find attribute {} in edge {}", - *attribute, edge_index - ))))?; + .remove_attribute(&attribute) + .map_err(PyMedRecordError::from)?; } Ok(())