use pyo3::exceptions::{PyIndexError, PyValueError};
use pyo3::prelude::*;
use nohash_hasher::{IntMap, IntSet};
use crate::id_tree::{IdTree, NodeId};
#[pyclass(name = "IDTree", unsendable)]
pub struct PyIdTree {
inner: IdTree,
}
#[pymethods]
impl PyIdTree {
#[new]
fn py_new(adj: std::collections::HashMap<usize, Vec<usize>>) -> PyResult<Self> {
let adj_map: IntMap<NodeId, IntSet<NodeId>> = adj
.into_iter()
.map(|(k, v)| (k, IntSet::from_iter(v)))
.collect();
Ok(Self {
inner: IdTree::from_adj_map(&adj_map),
})
}
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
}
}
#[pyo3(name = "insert_edge")]
fn py_insert_edge(&mut self, u: usize, v: usize) -> PyResult<i32> {
Ok(self.inner.insert_edge(u, v))
}
fn py_delete_edge(&mut self, u: usize, v: usize) -> PyResult<i32> {
Ok(self.inner.delete_edge(u, v))
}
#[pyo3(name = "query")]
fn py_query(&self, u: usize, v: usize) -> PyResult<bool> {
Ok(self.inner.query(u, v))
}
#[pyo3(name = "cycle_basis")]
fn py_cycle_basis(&mut self, root: Option<usize>) -> PyResult<Vec<Vec<usize>>> {
Ok(self.inner.cycle_basis(root))
}
#[pyo3(name = "node_connected_component")]
fn py_node_connected_component(&mut self, v: usize) -> PyResult<Vec<usize>> {
Ok(self.inner.node_connected_component(v))
}
#[pyo3(name = "num_connected_components")]
fn py_num_connected_components(&mut self) -> PyResult<usize> {
Ok(self.inner.num_connected_components())
}
#[pyo3(name = "connected_components")]
fn py_connected_components(&mut self) -> PyResult<Vec<Vec<usize>>> {
Ok(self.inner.connected_components())
}
#[pyo3(name = "active_nodes")]
fn py_active_nodes(&mut self) -> PyResult<Vec<usize>> {
Ok(self.inner.active_nodes_vec())
}
#[pyo3(name = "isolate_node")]
fn py_isolate_node(&mut self, v: usize) -> PyResult<()> {
self.inner.isolate_node(v);
Ok(())
}
#[pyo3(name = "isolate_nodes")]
fn py_isolate_nodes(&mut self, nodes: Vec<usize>) -> PyResult<()> {
self.inner.isolate_nodes(nodes);
Ok(())
}
#[pyo3(name = "is_isolated")]
fn py_is_isolated(&mut self, v: usize) -> PyResult<bool> {
Ok(self.inner.is_isolated(v))
}
#[pyo3(name = "degree")]
fn py_degree(&mut self, v: usize) -> PyResult<i32> {
Ok(self.inner.degree(v))
}
#[pyo3(name = "neighbors")]
fn py_neighbors(&mut self, v: usize) -> PyResult<Vec<usize>> {
Ok(self.inner.neighbors(v))
}
#[pyo3(name = "retain_active_nodes_from")]
fn py_retain_active_nodes_from(&mut self, from: Vec<usize>) -> PyResult<Vec<usize>> {
Ok(self.inner.retain_active_nodes_from(from))
}
#[pyo3(name = "shortest_path")]
fn py_shortest_path(&mut self, start: usize, target: usize) -> PyResult<Option<Vec<usize>>> {
Ok(self.inner.shortest_path(start, target))
}
}
#[pymodule]
pub fn python_idtree(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<PyIdTree>()?;
Ok(())
}