1#![allow(clippy::new_without_default, clippy::significant_drop_tightening)]
2
3pub mod attribute;
4mod borrowed;
5pub mod connector;
6pub mod datatype;
7pub mod errors;
8pub mod overview;
9pub mod plugins;
10pub mod querying;
11pub mod schema;
12pub mod traits;
13pub mod value;
14
15use crate::{
16 conversion_lut::ConversionLut,
17 graphrecord::{
18 overview::{PyGroupOverview, PyOverview},
19 plugins::PyPlugin,
20 },
21};
22use attribute::PyGraphRecordAttribute;
23use borrowed::BorrowedGraphRecord;
24use connector::PyConnector;
25use errors::PyGraphRecordError;
26use graphrecords_core::{
27 errors::GraphRecordError,
28 graphrecord::{
29 Attributes, EdgeIndex, GraphRecord, GraphRecordAttribute, GraphRecordValue,
30 connector::ConnectedGraphRecord, plugins::Plugin,
31 },
32};
33use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
34use pyo3::{
35 exceptions::PyRuntimeError,
36 prelude::*,
37 types::{PyBytes, PyDict, PyFunction},
38};
39use pyo3_polars::PyDataFrame;
40use querying::{PyReturnOperand, edges::PyEdgeOperand, nodes::PyNodeOperand};
41use schema::PySchema;
42use std::{
43 collections::HashMap,
44 ops::{Deref, DerefMut},
45 ptr::NonNull,
46};
47use traits::DeepInto;
48use value::PyGraphRecordValue;
49
50pub type PyAttributes = HashMap<PyGraphRecordAttribute, PyGraphRecordValue>;
51pub type PyGroup = PyGraphRecordAttribute;
52pub type PyPluginName = PyGraphRecordAttribute;
53pub type PyNodeIndex = PyGraphRecordAttribute;
54pub type PyEdgeIndex = EdgeIndex;
55type Lut<T> = ConversionLut<usize, fn(&Bound<'_, PyAny>) -> PyResult<T>>;
56
57#[pyclass(frozen)]
58#[derive(Debug)]
59pub struct PyGraphRecord {
60 inner: PyGraphRecordInner,
61}
62
63#[derive(Debug)]
64#[allow(clippy::large_enum_variant)]
65enum PyGraphRecordInner {
66 Owned(RwLock<GraphRecord>),
67 Connected(RwLock<ConnectedGraphRecord<PyConnector>>),
68 Borrowed(BorrowedGraphRecord),
69}
70
71pub(crate) enum InnerRef<'a> {
72 Owned(RwLockReadGuard<'a, GraphRecord>),
73 Connected(RwLockReadGuard<'a, ConnectedGraphRecord<PyConnector>>),
74 Borrowed(RwLockReadGuard<'a, Option<NonNull<GraphRecord>>>),
75}
76
77impl Deref for InnerRef<'_> {
78 type Target = GraphRecord;
79
80 fn deref(&self) -> &GraphRecord {
81 match self {
82 InnerRef::Owned(guard) => guard,
83 InnerRef::Connected(guard) => guard,
84 InnerRef::Borrowed(guard) => unsafe {
89 guard
90 .expect("Borrowed pointer must be Some when InnerRef is alive")
91 .as_ref()
92 },
93 }
94 }
95}
96
97pub(crate) enum InnerRefMut<'a> {
98 Owned(RwLockWriteGuard<'a, GraphRecord>),
99 Connected(RwLockWriteGuard<'a, ConnectedGraphRecord<PyConnector>>),
100 Borrowed(RwLockWriteGuard<'a, Option<NonNull<GraphRecord>>>),
101}
102
103impl Deref for InnerRefMut<'_> {
104 type Target = GraphRecord;
105
106 fn deref(&self) -> &GraphRecord {
107 match self {
108 InnerRefMut::Owned(guard) => guard,
109 InnerRefMut::Connected(guard) => guard,
110 InnerRefMut::Borrowed(guard) => unsafe {
114 guard
115 .expect("Borrowed pointer must be Some when InnerRefMut is alive")
116 .as_ref()
117 },
118 }
119 }
120}
121
122impl DerefMut for InnerRefMut<'_> {
123 fn deref_mut(&mut self) -> &mut GraphRecord {
124 match self {
125 InnerRefMut::Owned(guard) => &mut *guard,
126 InnerRefMut::Connected(guard) => &mut *guard,
127 InnerRefMut::Borrowed(guard) => unsafe {
133 guard
134 .expect("Borrowed pointer must be Some when InnerRefMut is alive")
135 .as_mut()
136 },
137 }
138 }
139}
140
141impl Clone for PyGraphRecord {
142 fn clone(&self) -> Self {
143 match &self.inner {
144 PyGraphRecordInner::Owned(lock) => Self {
145 inner: PyGraphRecordInner::Owned(RwLock::new(lock.read().clone())),
146 },
147 PyGraphRecordInner::Connected(lock) => Self {
148 inner: PyGraphRecordInner::Connected(RwLock::new(lock.read().clone())),
149 },
150 PyGraphRecordInner::Borrowed(_) => Self {
151 inner: PyGraphRecordInner::Borrowed(BorrowedGraphRecord::dead()),
152 },
153 }
154 }
155}
156
157impl PyGraphRecord {
158 pub(crate) fn inner(&self) -> PyResult<InnerRef<'_>> {
159 match &self.inner {
160 PyGraphRecordInner::Owned(lock) => Ok(InnerRef::Owned(lock.read())),
161 PyGraphRecordInner::Connected(lock) => Ok(InnerRef::Connected(lock.read())),
162 PyGraphRecordInner::Borrowed(borrowed) => {
163 let guard = borrowed.read();
164 if guard.is_some() {
165 Ok(InnerRef::Borrowed(guard))
166 } else {
167 Err(PyRuntimeError::new_err(
168 "GraphRecord reference is no longer valid (used outside callback scope)",
169 ))
170 }
171 }
172 }
173 }
174
175 pub(crate) fn connected(
176 &self,
177 ) -> PyResult<RwLockWriteGuard<'_, ConnectedGraphRecord<PyConnector>>> {
178 match &self.inner {
179 PyGraphRecordInner::Connected(lock) => Ok(lock.write()),
180 _ => Err(PyRuntimeError::new_err(
181 "GraphRecord has no connector attached",
182 )),
183 }
184 }
185
186 pub(crate) fn inner_mut(&self) -> PyResult<InnerRefMut<'_>> {
187 match &self.inner {
188 PyGraphRecordInner::Owned(lock) => Ok(InnerRefMut::Owned(lock.write())),
189 PyGraphRecordInner::Connected(lock) => Ok(InnerRefMut::Connected(lock.write())),
190 PyGraphRecordInner::Borrowed(borrowed) => {
191 if !borrowed.is_mutable() {
192 return Err(PyRuntimeError::new_err("GraphRecord is read-only"));
193 }
194 let guard = borrowed.write();
195 if guard.is_some() {
196 Ok(InnerRefMut::Borrowed(guard))
197 } else {
198 Err(PyRuntimeError::new_err(
199 "GraphRecord reference is no longer valid (used outside callback scope)",
200 ))
201 }
202 }
203 }
204 }
205}
206
207impl From<GraphRecord> for PyGraphRecord {
208 fn from(value: GraphRecord) -> Self {
209 Self {
210 inner: PyGraphRecordInner::Owned(RwLock::new(value)),
211 }
212 }
213}
214
215impl From<ConnectedGraphRecord<PyConnector>> for PyGraphRecord {
216 fn from(value: ConnectedGraphRecord<PyConnector>) -> Self {
217 Self {
218 inner: PyGraphRecordInner::Connected(RwLock::new(value)),
219 }
220 }
221}
222
223impl TryFrom<PyGraphRecord> for GraphRecord {
224 type Error = PyErr;
225
226 fn try_from(value: PyGraphRecord) -> PyResult<Self> {
227 match value.inner {
228 PyGraphRecordInner::Owned(lock) => Ok(lock.into_inner()),
229 PyGraphRecordInner::Connected(lock) => Ok(lock.into_inner().into()),
230 PyGraphRecordInner::Borrowed(_) => Err(PyRuntimeError::new_err(
231 "Cannot convert a borrowed PyGraphRecord into an owned GraphRecord",
232 )),
233 }
234 }
235}
236
237#[pymethods]
238impl PyGraphRecord {
239 #[new]
240 pub fn new() -> Self {
241 GraphRecord::new().into()
242 }
243
244 pub fn _to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
245 let bytes = bincode::serialize(&*self.inner()?)
246 .map_err(|_| {
247 GraphRecordError::ConversionError("Could not serialize GraphRecord".into())
248 })
249 .map_err(PyGraphRecordError::from)?;
250
251 Ok(PyBytes::new(py, &bytes))
252 }
253
254 #[staticmethod]
255 pub fn _from_bytes(data: &Bound<'_, PyBytes>) -> PyResult<Self> {
256 let graphrecord: GraphRecord = bincode::deserialize(data.as_bytes())
257 .map_err(|_| {
258 GraphRecordError::ConversionError("Could not deserialize GraphRecord".into())
259 })
260 .map_err(PyGraphRecordError::from)?;
261
262 Ok(graphrecord.into())
263 }
264
265 #[staticmethod]
266 pub fn with_schema(schema: PySchema) -> Self {
267 GraphRecord::with_schema(schema.into()).into()
268 }
269
270 #[staticmethod]
271 pub fn with_plugins(plugins: HashMap<PyPluginName, Py<PyAny>>) -> PyResult<Self> {
272 let plugins = plugins
273 .into_iter()
274 .map(|(name, plugin)| {
275 (
276 name.into(),
277 Box::new(PyPlugin::new(plugin)) as Box<dyn Plugin>,
278 )
279 })
280 .collect();
281
282 let graphrecord = GraphRecord::with_plugins(plugins).map_err(PyGraphRecordError::from)?;
283
284 Ok(graphrecord.into())
285 }
286
287 #[staticmethod]
288 #[pyo3(signature = (nodes, edges=None, schema=None))]
289 pub fn from_tuples(
290 nodes: Vec<(PyNodeIndex, PyAttributes)>,
291 edges: Option<Vec<(PyNodeIndex, PyNodeIndex, PyAttributes)>>,
292 schema: Option<PySchema>,
293 ) -> PyResult<Self> {
294 Ok(
295 GraphRecord::from_tuples(nodes.deep_into(), edges.deep_into(), schema.map(Into::into))
296 .map_err(PyGraphRecordError::from)?
297 .into(),
298 )
299 }
300
301 #[staticmethod]
302 #[pyo3(signature = (nodes_dataframes, edges_dataframes, schema=None))]
303 pub fn from_dataframes(
304 nodes_dataframes: Vec<(PyDataFrame, String)>,
305 edges_dataframes: Vec<(PyDataFrame, String, String)>,
306 schema: Option<PySchema>,
307 ) -> PyResult<Self> {
308 Ok(
309 GraphRecord::from_dataframes(
310 nodes_dataframes,
311 edges_dataframes,
312 schema.map(Into::into),
313 )
314 .map_err(PyGraphRecordError::from)?
315 .into(),
316 )
317 }
318
319 #[staticmethod]
320 #[pyo3(signature = (nodes_dataframes, schema=None))]
321 pub fn from_nodes_dataframes(
322 nodes_dataframes: Vec<(PyDataFrame, String)>,
323 schema: Option<PySchema>,
324 ) -> PyResult<Self> {
325 Ok(
326 GraphRecord::from_nodes_dataframes(nodes_dataframes, schema.map(Into::into))
327 .map_err(PyGraphRecordError::from)?
328 .into(),
329 )
330 }
331
332 #[staticmethod]
333 pub fn from_ron(path: &str) -> PyResult<Self> {
334 Ok(GraphRecord::from_ron(path)
335 .map_err(PyGraphRecordError::from)?
336 .into())
337 }
338
339 #[staticmethod]
340 pub fn with_connector(connector: Py<PyAny>) -> PyResult<Self> {
341 let connected = ConnectedGraphRecord::new(PyConnector::new(connector))
342 .map_err(PyGraphRecordError::from)?;
343
344 Ok(connected.into())
345 }
346
347 pub fn to_ron(&self, path: &str) -> PyResult<()> {
348 Ok(self
349 .inner()?
350 .to_ron(path)
351 .map_err(PyGraphRecordError::from)?)
352 }
353
354 #[allow(clippy::missing_panics_doc, reason = "infallible")]
355 pub fn to_dataframes(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
356 let export = self
357 .inner()?
358 .to_dataframes()
359 .map_err(PyGraphRecordError::from)?;
360
361 let outer_dict = PyDict::new(py);
362 let inner_dict = PyDict::new(py);
363
364 for (group, group_export) in export.groups {
365 let group_dict = PyDict::new(py);
366
367 let nodes_df = PyDataFrame(group_export.nodes);
368 group_dict
369 .set_item("nodes", nodes_df)
370 .expect("Setting item must succeed");
371
372 let edges_df = PyDataFrame(group_export.edges);
373 group_dict
374 .set_item("edges", edges_df)
375 .expect("Setting item must succeed");
376
377 inner_dict
378 .set_item(PyGraphRecordAttribute::from(group), group_dict)
379 .expect("Setting item must succeed");
380 }
381
382 outer_dict
383 .set_item("groups", inner_dict)
384 .expect("Setting item must succeed");
385
386 let ungrouped_dict = PyDict::new(py);
387
388 let nodes_df = PyDataFrame(export.ungrouped.nodes);
389 ungrouped_dict
390 .set_item("nodes", nodes_df)
391 .expect("Setting item must succeed");
392
393 let edges_df = PyDataFrame(export.ungrouped.edges);
394 ungrouped_dict
395 .set_item("edges", edges_df)
396 .expect("Setting item must succeed");
397
398 outer_dict
399 .set_item("ungrouped", ungrouped_dict)
400 .expect("Setting item must succeed");
401
402 Ok(outer_dict.into())
403 }
404
405 pub fn disconnect(&self) -> PyResult<Self> {
406 let graphrecord = self
407 .connected()?
408 .clone()
409 .disconnect()
410 .map_err(PyGraphRecordError::from)?;
411
412 Ok(graphrecord.into())
413 }
414
415 pub fn ingest(&self, data: Py<PyAny>) -> PyResult<()> {
416 self.connected()?
417 .ingest(data)
418 .map_err(PyGraphRecordError::from)?;
419
420 Ok(())
421 }
422
423 pub fn export(&self) -> PyResult<Py<PyAny>> {
424 let data = self
425 .connected()?
426 .export()
427 .map_err(PyGraphRecordError::from)?;
428
429 Ok(data)
430 }
431
432 pub fn add_plugin(&self, name: PyPluginName, plugin: Py<PyAny>) -> PyResult<()> {
433 let mut graphrecord = self.inner_mut()?;
434
435 graphrecord
436 .add_plugin(name.into(), Box::new(PyPlugin::new(plugin)))
437 .map_err(PyGraphRecordError::from)?;
438
439 Ok(())
440 }
441
442 pub fn remove_plugin(&self, name: PyPluginName) -> PyResult<()> {
443 let mut graphrecord = self.inner_mut()?;
444
445 graphrecord
446 .remove_plugin(&name.into())
447 .map_err(PyGraphRecordError::from)?;
448
449 Ok(())
450 }
451
452 #[getter]
453 pub fn plugins(&self) -> PyResult<Vec<PyPluginName>> {
454 Ok(self
455 .inner()?
456 .plugin_names()
457 .cloned()
458 .map(std::convert::Into::into)
459 .collect())
460 }
461
462 pub fn get_schema(&self) -> PyResult<PySchema> {
463 Ok(self.inner()?.get_schema().clone().into())
464 }
465
466 #[pyo3(signature = (schema, bypass_plugins=false))]
467 pub fn set_schema(&self, schema: PySchema, bypass_plugins: bool) -> PyResult<()> {
468 let mut graphrecord = self.inner_mut()?;
469
470 if bypass_plugins {
471 Ok(graphrecord
472 .set_schema_bypass_plugins(schema.into())
473 .map_err(PyGraphRecordError::from)?)
474 } else {
475 Ok(graphrecord
476 .set_schema(schema.into())
477 .map_err(PyGraphRecordError::from)?)
478 }
479 }
480
481 #[pyo3(signature = (bypass_plugins=false))]
482 pub fn freeze_schema(&self, bypass_plugins: bool) -> PyResult<()> {
483 let mut graphrecord = self.inner_mut()?;
484
485 if bypass_plugins {
486 Ok(graphrecord
487 .freeze_schema_bypass_plugins()
488 .map_err(PyGraphRecordError::from)?)
489 } else {
490 Ok(graphrecord
491 .freeze_schema()
492 .map_err(PyGraphRecordError::from)?)
493 }
494 }
495
496 #[pyo3(signature = (bypass_plugins=false))]
497 pub fn unfreeze_schema(&self, bypass_plugins: bool) -> PyResult<()> {
498 let mut graphrecord = self.inner_mut()?;
499
500 if bypass_plugins {
501 Ok(graphrecord
502 .unfreeze_schema_bypass_plugins()
503 .map_err(PyGraphRecordError::from)?)
504 } else {
505 Ok(graphrecord
506 .unfreeze_schema()
507 .map_err(PyGraphRecordError::from)?)
508 }
509 }
510
511 #[getter]
512 pub fn nodes(&self) -> PyResult<Vec<PyNodeIndex>> {
513 Ok(self
514 .inner()?
515 .node_indices()
516 .map(|node_index| node_index.clone().into())
517 .collect())
518 }
519
520 pub fn node(
521 &self,
522 node_index: Vec<PyNodeIndex>,
523 ) -> PyResult<HashMap<PyNodeIndex, PyAttributes>> {
524 let graphrecord = self.inner()?;
525
526 node_index
527 .into_iter()
528 .map(|node_index| {
529 let node_attributes = graphrecord
530 .node_attributes(&node_index)
531 .map_err(PyGraphRecordError::from)?;
532
533 Ok((node_index, node_attributes.deep_into()))
534 })
535 .collect()
536 }
537
538 #[getter]
539 pub fn edges(&self) -> PyResult<Vec<EdgeIndex>> {
540 Ok(self.inner()?.edge_indices().copied().collect())
541 }
542
543 pub fn edge(&self, edge_index: Vec<EdgeIndex>) -> PyResult<HashMap<EdgeIndex, PyAttributes>> {
544 let graphrecord = self.inner()?;
545
546 edge_index
547 .into_iter()
548 .map(|edge_index| {
549 let edge_attributes = graphrecord
550 .edge_attributes(&edge_index)
551 .map_err(PyGraphRecordError::from)?;
552
553 Ok((edge_index, edge_attributes.deep_into()))
554 })
555 .collect()
556 }
557
558 #[getter]
559 pub fn groups(&self) -> PyResult<Vec<PyGroup>> {
560 Ok(self
561 .inner()?
562 .groups()
563 .map(|group| group.clone().into())
564 .collect())
565 }
566
567 pub fn outgoing_edges(
568 &self,
569 node_index: Vec<PyNodeIndex>,
570 ) -> PyResult<HashMap<PyNodeIndex, Vec<EdgeIndex>>> {
571 let graphrecord = self.inner()?;
572
573 node_index
574 .into_iter()
575 .map(|node_index| {
576 let edges = graphrecord
577 .outgoing_edges(&node_index)
578 .map_err(PyGraphRecordError::from)?
579 .copied()
580 .collect();
581
582 Ok((node_index, edges))
583 })
584 .collect()
585 }
586
587 pub fn incoming_edges(
588 &self,
589 node_index: Vec<PyNodeIndex>,
590 ) -> PyResult<HashMap<PyNodeIndex, Vec<EdgeIndex>>> {
591 let graphrecord = self.inner()?;
592
593 node_index
594 .into_iter()
595 .map(|node_index| {
596 let edges = graphrecord
597 .incoming_edges(&node_index)
598 .map_err(PyGraphRecordError::from)?
599 .copied()
600 .collect();
601
602 Ok((node_index, edges))
603 })
604 .collect()
605 }
606
607 pub fn edge_endpoints(
608 &self,
609 edge_index: Vec<EdgeIndex>,
610 ) -> PyResult<HashMap<EdgeIndex, (PyNodeIndex, PyNodeIndex)>> {
611 let graphrecord = self.inner()?;
612
613 edge_index
614 .into_iter()
615 .map(|edge_index| {
616 let edge_endpoints = graphrecord
617 .edge_endpoints(&edge_index)
618 .map_err(PyGraphRecordError::from)?;
619
620 Ok((
621 edge_index,
622 (
623 edge_endpoints.0.clone().into(),
624 edge_endpoints.1.clone().into(),
625 ),
626 ))
627 })
628 .collect()
629 }
630
631 pub fn edges_connecting(
632 &self,
633 source_node_indices: Vec<PyNodeIndex>,
634 target_node_indices: Vec<PyNodeIndex>,
635 ) -> PyResult<Vec<EdgeIndex>> {
636 let source_node_indices: Vec<GraphRecordAttribute> = source_node_indices.deep_into();
637 let target_node_indices: Vec<GraphRecordAttribute> = target_node_indices.deep_into();
638
639 Ok(self
640 .inner()?
641 .edges_connecting(
642 source_node_indices.iter().collect(),
643 target_node_indices.iter().collect(),
644 )
645 .copied()
646 .collect())
647 }
648
649 pub fn edges_connecting_undirected(
650 &self,
651 first_node_indices: Vec<PyNodeIndex>,
652 second_node_indices: Vec<PyNodeIndex>,
653 ) -> PyResult<Vec<EdgeIndex>> {
654 let first_node_indices: Vec<GraphRecordAttribute> = first_node_indices.deep_into();
655 let second_node_indices: Vec<GraphRecordAttribute> = second_node_indices.deep_into();
656
657 Ok(self
658 .inner()?
659 .edges_connecting_undirected(
660 first_node_indices.iter().collect(),
661 second_node_indices.iter().collect(),
662 )
663 .copied()
664 .collect())
665 }
666
667 #[pyo3(signature = (node_indices, bypass_plugins=false))]
668 pub fn remove_nodes(
669 &self,
670 node_indices: Vec<PyNodeIndex>,
671 bypass_plugins: bool,
672 ) -> PyResult<HashMap<PyNodeIndex, PyAttributes>> {
673 let mut graphrecord = self.inner_mut()?;
674
675 if bypass_plugins {
676 node_indices
677 .into_iter()
678 .map(|node_index| {
679 let attributes = graphrecord
680 .remove_node_bypass_plugins(&node_index)
681 .map_err(PyGraphRecordError::from)?;
682 Ok((node_index, attributes.deep_into()))
683 })
684 .collect()
685 } else {
686 node_indices
687 .into_iter()
688 .map(|node_index| {
689 let attributes = graphrecord
690 .remove_node(&node_index)
691 .map_err(PyGraphRecordError::from)?;
692 Ok((node_index, attributes.deep_into()))
693 })
694 .collect()
695 }
696 }
697
698 pub fn replace_node_attributes(
699 &self,
700 node_indices: Vec<PyNodeIndex>,
701 attributes: PyAttributes,
702 ) -> PyResult<()> {
703 let mut graphrecord = self.inner_mut()?;
704
705 let attributes: Attributes = attributes.deep_into();
706
707 for node_index in node_indices {
708 let mut current_attributes = graphrecord
709 .node_attributes_mut(&node_index)
710 .map_err(PyGraphRecordError::from)?;
711
712 current_attributes
713 .replace_attributes(attributes.clone())
714 .map_err(PyGraphRecordError::from)?;
715 }
716
717 Ok(())
718 }
719
720 pub fn update_node_attribute(
721 &self,
722 node_indices: Vec<PyNodeIndex>,
723 attribute: PyGraphRecordAttribute,
724 value: PyGraphRecordValue,
725 ) -> PyResult<()> {
726 let mut graphrecord = self.inner_mut()?;
727
728 let attribute: GraphRecordAttribute = attribute.into();
729 let value: GraphRecordValue = value.into();
730
731 for node_index in node_indices {
732 let mut node_attributes = graphrecord
733 .node_attributes_mut(&node_index)
734 .map_err(PyGraphRecordError::from)?;
735
736 node_attributes
737 .update_attribute(&attribute, value.clone())
738 .map_err(PyGraphRecordError::from)?;
739 }
740
741 Ok(())
742 }
743
744 pub fn remove_node_attribute(
745 &self,
746 node_indices: Vec<PyNodeIndex>,
747 attribute: PyGraphRecordAttribute,
748 ) -> PyResult<()> {
749 let mut graphrecord = self.inner_mut()?;
750
751 let attribute: GraphRecordAttribute = attribute.into();
752
753 for node_index in node_indices {
754 let mut node_attributes = graphrecord
755 .node_attributes_mut(&node_index)
756 .map_err(PyGraphRecordError::from)?;
757
758 node_attributes
759 .remove_attribute(&attribute)
760 .map_err(PyGraphRecordError::from)?;
761 }
762
763 Ok(())
764 }
765
766 #[pyo3(signature = (nodes, bypass_plugins=false))]
767 pub fn add_nodes(
768 &self,
769 nodes: Vec<(PyNodeIndex, PyAttributes)>,
770 bypass_plugins: bool,
771 ) -> PyResult<()> {
772 let mut graphrecord = self.inner_mut()?;
773
774 if bypass_plugins {
775 Ok(graphrecord
776 .add_nodes_bypass_plugins(nodes.deep_into())
777 .map_err(PyGraphRecordError::from)?)
778 } else {
779 Ok(graphrecord
780 .add_nodes(nodes.deep_into())
781 .map_err(PyGraphRecordError::from)?)
782 }
783 }
784
785 #[pyo3(signature = (nodes, group, bypass_plugins=false))]
786 pub fn add_nodes_with_group(
787 &self,
788 nodes: Vec<(PyNodeIndex, PyAttributes)>,
789 group: PyGroup,
790 bypass_plugins: bool,
791 ) -> PyResult<()> {
792 let mut graphrecord = self.inner_mut()?;
793
794 if bypass_plugins {
795 Ok(graphrecord
796 .add_nodes_with_group_bypass_plugins(nodes.deep_into(), group.into())
797 .map_err(PyGraphRecordError::from)?)
798 } else {
799 Ok(graphrecord
800 .add_nodes_with_group(nodes.deep_into(), group.into())
801 .map_err(PyGraphRecordError::from)?)
802 }
803 }
804
805 #[pyo3(signature = (nodes_dataframes, bypass_plugins=false))]
806 pub fn add_nodes_dataframes(
807 &self,
808 nodes_dataframes: Vec<(PyDataFrame, String)>,
809 bypass_plugins: bool,
810 ) -> PyResult<()> {
811 let mut graphrecord = self.inner_mut()?;
812
813 if bypass_plugins {
814 Ok(graphrecord
815 .add_nodes_dataframes_bypass_plugins(nodes_dataframes)
816 .map_err(PyGraphRecordError::from)?)
817 } else {
818 Ok(graphrecord
819 .add_nodes_dataframes(nodes_dataframes)
820 .map_err(PyGraphRecordError::from)?)
821 }
822 }
823
824 #[pyo3(signature = (nodes_dataframes, group, bypass_plugins=false))]
825 pub fn add_nodes_dataframes_with_group(
826 &self,
827 nodes_dataframes: Vec<(PyDataFrame, String)>,
828 group: PyGroup,
829 bypass_plugins: bool,
830 ) -> PyResult<()> {
831 let mut graphrecord = self.inner_mut()?;
832
833 if bypass_plugins {
834 Ok(graphrecord
835 .add_nodes_dataframes_with_group_bypass_plugins(nodes_dataframes, group.into())
836 .map_err(PyGraphRecordError::from)?)
837 } else {
838 Ok(graphrecord
839 .add_nodes_dataframes_with_group(nodes_dataframes, group.into())
840 .map_err(PyGraphRecordError::from)?)
841 }
842 }
843
844 #[pyo3(signature = (edge_indices, bypass_plugins=false))]
845 pub fn remove_edges(
846 &self,
847 edge_indices: Vec<EdgeIndex>,
848 bypass_plugins: bool,
849 ) -> PyResult<HashMap<EdgeIndex, PyAttributes>> {
850 let mut graphrecord = self.inner_mut()?;
851
852 if bypass_plugins {
853 edge_indices
854 .into_iter()
855 .map(|edge_index| {
856 let attributes = graphrecord
857 .remove_edge_bypass_plugins(&edge_index)
858 .map_err(PyGraphRecordError::from)?;
859 Ok((edge_index, attributes.deep_into()))
860 })
861 .collect()
862 } else {
863 edge_indices
864 .into_iter()
865 .map(|edge_index| {
866 let attributes = graphrecord
867 .remove_edge(&edge_index)
868 .map_err(PyGraphRecordError::from)?;
869 Ok((edge_index, attributes.deep_into()))
870 })
871 .collect()
872 }
873 }
874
875 pub fn replace_edge_attributes(
876 &self,
877 edge_indices: Vec<EdgeIndex>,
878 attributes: PyAttributes,
879 ) -> PyResult<()> {
880 let mut graphrecord = self.inner_mut()?;
881
882 let attributes: Attributes = attributes.deep_into();
883
884 for edge_index in edge_indices {
885 let mut current_attributes = graphrecord
886 .edge_attributes_mut(&edge_index)
887 .map_err(PyGraphRecordError::from)?;
888
889 current_attributes
890 .replace_attributes(attributes.clone())
891 .map_err(PyGraphRecordError::from)?;
892 }
893
894 Ok(())
895 }
896
897 pub fn update_edge_attribute(
898 &self,
899 edge_indices: Vec<EdgeIndex>,
900 attribute: PyGraphRecordAttribute,
901 value: PyGraphRecordValue,
902 ) -> PyResult<()> {
903 let mut graphrecord = self.inner_mut()?;
904
905 let attribute: GraphRecordAttribute = attribute.into();
906 let value: GraphRecordValue = value.into();
907
908 for edge_index in edge_indices {
909 let mut edge_attributes = graphrecord
910 .edge_attributes_mut(&edge_index)
911 .map_err(PyGraphRecordError::from)?;
912
913 edge_attributes
914 .update_attribute(&attribute, value.clone())
915 .map_err(PyGraphRecordError::from)?;
916 }
917
918 Ok(())
919 }
920
921 pub fn remove_edge_attribute(
922 &self,
923 edge_indices: Vec<EdgeIndex>,
924 attribute: PyGraphRecordAttribute,
925 ) -> PyResult<()> {
926 let mut graphrecord = self.inner_mut()?;
927
928 let attribute: GraphRecordAttribute = attribute.into();
929
930 for edge_index in edge_indices {
931 let mut edge_attributes = graphrecord
932 .edge_attributes_mut(&edge_index)
933 .map_err(PyGraphRecordError::from)?;
934
935 edge_attributes
936 .remove_attribute(&attribute)
937 .map_err(PyGraphRecordError::from)?;
938 }
939
940 Ok(())
941 }
942
943 #[pyo3(signature = (relations, bypass_plugins=false))]
944 pub fn add_edges(
945 &self,
946 relations: Vec<(PyNodeIndex, PyNodeIndex, PyAttributes)>,
947 bypass_plugins: bool,
948 ) -> PyResult<Vec<EdgeIndex>> {
949 let mut graphrecord = self.inner_mut()?;
950
951 if bypass_plugins {
952 Ok(graphrecord
953 .add_edges_bypass_plugins(relations.deep_into())
954 .map_err(PyGraphRecordError::from)?)
955 } else {
956 Ok(graphrecord
957 .add_edges(relations.deep_into())
958 .map_err(PyGraphRecordError::from)?)
959 }
960 }
961
962 #[pyo3(signature = (relations, group, bypass_plugins=false))]
963 pub fn add_edges_with_group(
964 &self,
965 relations: Vec<(PyNodeIndex, PyNodeIndex, PyAttributes)>,
966 group: PyGroup,
967 bypass_plugins: bool,
968 ) -> PyResult<Vec<EdgeIndex>> {
969 let mut graphrecord = self.inner_mut()?;
970
971 if bypass_plugins {
972 Ok(graphrecord
973 .add_edges_with_group_bypass_plugins(relations.deep_into(), &group)
974 .map_err(PyGraphRecordError::from)?)
975 } else {
976 Ok(graphrecord
977 .add_edges_with_group(relations.deep_into(), &group)
978 .map_err(PyGraphRecordError::from)?)
979 }
980 }
981
982 #[pyo3(signature = (edges_dataframes, bypass_plugins=false))]
983 pub fn add_edges_dataframes(
984 &self,
985 edges_dataframes: Vec<(PyDataFrame, String, String)>,
986 bypass_plugins: bool,
987 ) -> PyResult<Vec<EdgeIndex>> {
988 let mut graphrecord = self.inner_mut()?;
989
990 if bypass_plugins {
991 Ok(graphrecord
992 .add_edges_dataframes_bypass_plugins(edges_dataframes)
993 .map_err(PyGraphRecordError::from)?)
994 } else {
995 Ok(graphrecord
996 .add_edges_dataframes(edges_dataframes)
997 .map_err(PyGraphRecordError::from)?)
998 }
999 }
1000
1001 #[pyo3(signature = (edges_dataframes, group, bypass_plugins=false))]
1002 pub fn add_edges_dataframes_with_group(
1003 &self,
1004 edges_dataframes: Vec<(PyDataFrame, String, String)>,
1005 group: PyGroup,
1006 bypass_plugins: bool,
1007 ) -> PyResult<Vec<EdgeIndex>> {
1008 let mut graphrecord = self.inner_mut()?;
1009
1010 if bypass_plugins {
1011 Ok(graphrecord
1012 .add_edges_dataframes_with_group_bypass_plugins(edges_dataframes, &group)
1013 .map_err(PyGraphRecordError::from)?)
1014 } else {
1015 Ok(graphrecord
1016 .add_edges_dataframes_with_group(edges_dataframes, &group)
1017 .map_err(PyGraphRecordError::from)?)
1018 }
1019 }
1020
1021 #[pyo3(signature = (group, node_indices_to_add=None, edge_indices_to_add=None, bypass_plugins=false))]
1022 pub fn add_group(
1023 &self,
1024 group: PyGroup,
1025 node_indices_to_add: Option<Vec<PyNodeIndex>>,
1026 edge_indices_to_add: Option<Vec<EdgeIndex>>,
1027 bypass_plugins: bool,
1028 ) -> PyResult<()> {
1029 let mut graphrecord = self.inner_mut()?;
1030
1031 if bypass_plugins {
1032 Ok(graphrecord
1033 .add_group_bypass_plugins(
1034 group.into(),
1035 node_indices_to_add.deep_into(),
1036 edge_indices_to_add,
1037 )
1038 .map_err(PyGraphRecordError::from)?)
1039 } else {
1040 Ok(graphrecord
1041 .add_group(
1042 group.into(),
1043 node_indices_to_add.deep_into(),
1044 edge_indices_to_add,
1045 )
1046 .map_err(PyGraphRecordError::from)?)
1047 }
1048 }
1049
1050 #[pyo3(signature = (group, bypass_plugins=false))]
1051 pub fn remove_groups(&self, group: Vec<PyGroup>, bypass_plugins: bool) -> PyResult<()> {
1052 let mut graphrecord = self.inner_mut()?;
1053
1054 if bypass_plugins {
1055 group.into_iter().try_for_each(|group| {
1056 graphrecord
1057 .remove_group_bypass_plugins(&group)
1058 .map_err(PyGraphRecordError::from)?;
1059 Ok(())
1060 })
1061 } else {
1062 group.into_iter().try_for_each(|group| {
1063 graphrecord
1064 .remove_group(&group)
1065 .map_err(PyGraphRecordError::from)?;
1066 Ok(())
1067 })
1068 }
1069 }
1070
1071 #[pyo3(signature = (group, node_index, bypass_plugins=false))]
1072 pub fn add_nodes_to_group(
1073 &self,
1074 group: PyGroup,
1075 node_index: Vec<PyNodeIndex>,
1076 bypass_plugins: bool,
1077 ) -> PyResult<()> {
1078 let mut graphrecord = self.inner_mut()?;
1079
1080 if bypass_plugins {
1081 node_index.into_iter().try_for_each(|node_index| {
1082 Ok(graphrecord
1083 .add_node_to_group_bypass_plugins(group.clone().into(), node_index.into())
1084 .map_err(PyGraphRecordError::from)?)
1085 })
1086 } else {
1087 node_index.into_iter().try_for_each(|node_index| {
1088 Ok(graphrecord
1089 .add_node_to_group(group.clone().into(), node_index.into())
1090 .map_err(PyGraphRecordError::from)?)
1091 })
1092 }
1093 }
1094
1095 #[pyo3(signature = (group, edge_index, bypass_plugins=false))]
1096 pub fn add_edges_to_group(
1097 &self,
1098 group: PyGroup,
1099 edge_index: Vec<EdgeIndex>,
1100 bypass_plugins: bool,
1101 ) -> PyResult<()> {
1102 let mut graphrecord = self.inner_mut()?;
1103
1104 if bypass_plugins {
1105 edge_index.into_iter().try_for_each(|edge_index| {
1106 Ok(graphrecord
1107 .add_edge_to_group_bypass_plugins(group.clone().into(), edge_index)
1108 .map_err(PyGraphRecordError::from)?)
1109 })
1110 } else {
1111 edge_index.into_iter().try_for_each(|edge_index| {
1112 Ok(graphrecord
1113 .add_edge_to_group(group.clone().into(), edge_index)
1114 .map_err(PyGraphRecordError::from)?)
1115 })
1116 }
1117 }
1118
1119 #[pyo3(signature = (group, node_index, bypass_plugins=false))]
1120 pub fn remove_nodes_from_group(
1121 &self,
1122 group: PyGroup,
1123 node_index: Vec<PyNodeIndex>,
1124 bypass_plugins: bool,
1125 ) -> PyResult<()> {
1126 let mut graphrecord = self.inner_mut()?;
1127
1128 if bypass_plugins {
1129 node_index.into_iter().try_for_each(|node_index| {
1130 Ok(graphrecord
1131 .remove_node_from_group_bypass_plugins(&group, &node_index)
1132 .map_err(PyGraphRecordError::from)?)
1133 })
1134 } else {
1135 node_index.into_iter().try_for_each(|node_index| {
1136 Ok(graphrecord
1137 .remove_node_from_group(&group, &node_index)
1138 .map_err(PyGraphRecordError::from)?)
1139 })
1140 }
1141 }
1142
1143 #[pyo3(signature = (group, edge_index, bypass_plugins=false))]
1144 pub fn remove_edges_from_group(
1145 &self,
1146 group: PyGroup,
1147 edge_index: Vec<EdgeIndex>,
1148 bypass_plugins: bool,
1149 ) -> PyResult<()> {
1150 let mut graphrecord = self.inner_mut()?;
1151
1152 if bypass_plugins {
1153 edge_index.into_iter().try_for_each(|edge_index| {
1154 Ok(graphrecord
1155 .remove_edge_from_group_bypass_plugins(&group, &edge_index)
1156 .map_err(PyGraphRecordError::from)?)
1157 })
1158 } else {
1159 edge_index.into_iter().try_for_each(|edge_index| {
1160 Ok(graphrecord
1161 .remove_edge_from_group(&group, &edge_index)
1162 .map_err(PyGraphRecordError::from)?)
1163 })
1164 }
1165 }
1166
1167 pub fn nodes_in_group(
1168 &self,
1169 group: Vec<PyGroup>,
1170 ) -> PyResult<HashMap<PyGroup, Vec<PyNodeIndex>>> {
1171 let graphrecord = self.inner()?;
1172
1173 group
1174 .into_iter()
1175 .map(|group| {
1176 let nodes_attributes = graphrecord
1177 .nodes_in_group(&group)
1178 .map_err(PyGraphRecordError::from)?
1179 .map(|node_index| node_index.clone().into())
1180 .collect();
1181
1182 Ok((group, nodes_attributes))
1183 })
1184 .collect()
1185 }
1186
1187 pub fn ungrouped_nodes(&self) -> PyResult<Vec<PyNodeIndex>> {
1188 Ok(self
1189 .inner()?
1190 .ungrouped_nodes()
1191 .map(|node_index| node_index.clone().into())
1192 .collect())
1193 }
1194
1195 pub fn edges_in_group(
1196 &self,
1197 group: Vec<PyGroup>,
1198 ) -> PyResult<HashMap<PyGroup, Vec<EdgeIndex>>> {
1199 let graphrecord = self.inner()?;
1200
1201 group
1202 .into_iter()
1203 .map(|group| {
1204 let edges = graphrecord
1205 .edges_in_group(&group)
1206 .map_err(PyGraphRecordError::from)?
1207 .copied()
1208 .collect();
1209
1210 Ok((group, edges))
1211 })
1212 .collect()
1213 }
1214
1215 pub fn ungrouped_edges(&self) -> PyResult<Vec<EdgeIndex>> {
1216 Ok(self.inner()?.ungrouped_edges().copied().collect())
1217 }
1218
1219 pub fn groups_of_node(
1220 &self,
1221 node_index: Vec<PyNodeIndex>,
1222 ) -> PyResult<HashMap<PyNodeIndex, Vec<PyGroup>>> {
1223 let graphrecord = self.inner()?;
1224
1225 node_index
1226 .into_iter()
1227 .map(|node_index| {
1228 let groups = graphrecord
1229 .groups_of_node(&node_index)
1230 .map_err(PyGraphRecordError::from)?
1231 .map(|node_index| node_index.clone().into())
1232 .collect();
1233
1234 Ok((node_index, groups))
1235 })
1236 .collect()
1237 }
1238
1239 pub fn groups_of_edge(
1240 &self,
1241 edge_index: Vec<EdgeIndex>,
1242 ) -> PyResult<HashMap<EdgeIndex, Vec<PyGroup>>> {
1243 let graphrecord = self.inner()?;
1244
1245 edge_index
1246 .into_iter()
1247 .map(|edge_index| {
1248 let groups = graphrecord
1249 .groups_of_edge(&edge_index)
1250 .map_err(PyGraphRecordError::from)?
1251 .map(|group| group.clone().into())
1252 .collect();
1253
1254 Ok((edge_index, groups))
1255 })
1256 .collect()
1257 }
1258
1259 pub fn node_count(&self) -> PyResult<usize> {
1260 Ok(self.inner()?.node_count())
1261 }
1262
1263 pub fn edge_count(&self) -> PyResult<usize> {
1264 Ok(self.inner()?.edge_count())
1265 }
1266
1267 pub fn group_count(&self) -> PyResult<usize> {
1268 Ok(self.inner()?.group_count())
1269 }
1270
1271 pub fn contains_node(&self, node_index: PyNodeIndex) -> PyResult<bool> {
1272 Ok(self.inner()?.contains_node(&node_index.into()))
1273 }
1274
1275 pub fn contains_edge(&self, edge_index: EdgeIndex) -> PyResult<bool> {
1276 Ok(self.inner()?.contains_edge(&edge_index))
1277 }
1278
1279 pub fn contains_group(&self, group: PyGroup) -> PyResult<bool> {
1280 Ok(self.inner()?.contains_group(&group.into()))
1281 }
1282
1283 pub fn neighbors_outgoing(
1284 &self,
1285 node_indices: Vec<PyNodeIndex>,
1286 ) -> PyResult<HashMap<PyNodeIndex, Vec<PyNodeIndex>>> {
1287 let graphrecord = self.inner()?;
1288
1289 node_indices
1290 .into_iter()
1291 .map(|node_index| {
1292 let neighbors = graphrecord
1293 .neighbors_outgoing(&node_index)
1294 .map_err(PyGraphRecordError::from)?
1295 .map(|neighbor| neighbor.clone().into())
1296 .collect();
1297
1298 Ok((node_index, neighbors))
1299 })
1300 .collect()
1301 }
1302
1303 pub fn neighbors_incoming(
1304 &self,
1305 node_indices: Vec<PyNodeIndex>,
1306 ) -> PyResult<HashMap<PyNodeIndex, Vec<PyNodeIndex>>> {
1307 let graphrecord = self.inner()?;
1308
1309 node_indices
1310 .into_iter()
1311 .map(|node_index| {
1312 let neighbors = graphrecord
1313 .neighbors_incoming(&node_index)
1314 .map_err(PyGraphRecordError::from)?
1315 .map(|neighbor| neighbor.clone().into())
1316 .collect();
1317
1318 Ok((node_index, neighbors))
1319 })
1320 .collect()
1321 }
1322
1323 pub fn neighbors_undirected(
1324 &self,
1325 node_indices: Vec<PyNodeIndex>,
1326 ) -> PyResult<HashMap<PyNodeIndex, Vec<PyNodeIndex>>> {
1327 let graphrecord = self.inner()?;
1328
1329 node_indices
1330 .into_iter()
1331 .map(|node_index| {
1332 let neighbors = graphrecord
1333 .neighbors_undirected(&node_index)
1334 .map_err(PyGraphRecordError::from)?
1335 .map(|neighbor| neighbor.clone().into())
1336 .collect();
1337
1338 Ok((node_index, neighbors))
1339 })
1340 .collect()
1341 }
1342
1343 #[pyo3(signature = (bypass_plugins=false))]
1344 pub fn clear(&self, bypass_plugins: bool) -> PyResult<()> {
1345 let mut graphrecord = self.inner_mut()?;
1346
1347 if bypass_plugins {
1348 Ok(graphrecord
1349 .clear_bypass_plugins()
1350 .map_err(PyGraphRecordError::from)?)
1351 } else {
1352 Ok(graphrecord.clear().map_err(PyGraphRecordError::from)?)
1353 }
1354 }
1355
1356 pub fn query_nodes(
1360 &self,
1361 py: Python<'_>,
1362 query: &Bound<'_, PyFunction>,
1363 ) -> PyResult<Py<PyAny>> {
1364 let graphrecord = self.inner()?;
1365
1366 let result = graphrecord
1367 .query_nodes(|nodes| {
1368 let result = query
1369 .call1((PyNodeOperand::from(nodes.clone()),))
1370 .expect("Call should succeed");
1371
1372 result
1373 .extract::<PyReturnOperand>()
1374 .expect("Extraction must succeed")
1375 })
1376 .evaluate()
1377 .map_err(PyGraphRecordError::from)?;
1378
1379 Ok(result.into_pyobject(py)?.unbind())
1380 }
1381
1382 pub fn query_edges(
1386 &self,
1387 py: Python<'_>,
1388 query: &Bound<'_, PyFunction>,
1389 ) -> PyResult<Py<PyAny>> {
1390 let graphrecord = self.inner()?;
1391
1392 let result = graphrecord
1393 .query_edges(|edges| {
1394 let result = query
1395 .call1((PyEdgeOperand::from(edges.clone()),))
1396 .expect("Call should succeed");
1397
1398 result
1399 .extract::<PyReturnOperand>()
1400 .expect("Extraction must succeed")
1401 })
1402 .evaluate()
1403 .map_err(PyGraphRecordError::from)?;
1404
1405 Ok(result.into_pyobject(py)?.unbind())
1406 }
1407
1408 #[allow(clippy::should_implement_trait)]
1409 pub fn clone(&self) -> Self {
1410 Clone::clone(self)
1411 }
1412
1413 pub fn overview(&self, truncate_details: Option<usize>) -> PyResult<PyOverview> {
1414 Ok(self
1415 .inner()?
1416 .overview(truncate_details)
1417 .map_err(PyGraphRecordError::from)?
1418 .into())
1419 }
1420
1421 pub fn group_overview(
1422 &self,
1423 group: PyGroup,
1424 truncate_details: Option<usize>,
1425 ) -> PyResult<PyGroupOverview> {
1426 Ok(self
1427 .inner()?
1428 .group_overview(&group.into(), truncate_details)
1429 .map_err(PyGraphRecordError::from)?
1430 .into())
1431 }
1432}