Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow some graph and digraph methods to take iterables/generators #1292

Merged
merged 6 commits into from
Oct 11, 2024
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 14 additions & 14 deletions rustworkx/rustworkx.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -1153,14 +1153,14 @@ class PyGraph(Generic[_S, _T]):
def add_edge(self, node_a: int, node_b: int, edge: _T, /) -> int: ...
def add_edges_from(
self,
obj_list: Sequence[tuple[int, int, _T]],
obj_list: Iterable[tuple[int, int, _T]],
/,
) -> list[int]: ...
def add_edges_from_no_data(
self: PyGraph[_S, _T | None], obj_list: Sequence[tuple[int, int]], /
self: PyGraph[_S, _T | None], obj_list: Iterable[tuple[int, int]], /
) -> list[int]: ...
def add_node(self, obj: _S, /) -> int: ...
def add_nodes_from(self, obj_list: Sequence[_S], /) -> NodeIndices: ...
def add_nodes_from(self, obj_list: Iterable[_S], /) -> NodeIndices: ...
def adj(self, node: int, /) -> dict[int, _T]: ...
def clear(self) -> None: ...
def clear_edges(self) -> None: ...
Expand Down Expand Up @@ -1188,11 +1188,11 @@ class PyGraph(Generic[_S, _T]):
def edges(self) -> list[_T]: ...
def edge_subgraph(self, edge_list: Sequence[tuple[int, int]], /) -> PyGraph[_S, _T]: ...
def extend_from_edge_list(
self: PyGraph[_S | None, _T | None], edge_list: Sequence[tuple[int, int]], /
self: PyGraph[_S | None, _T | None], edge_list: Iterable[tuple[int, int]], /
) -> None: ...
def extend_from_weighted_edge_list(
self: PyGraph[_S | None, _T],
edge_list: Sequence[tuple[int, int, _T]],
edge_list: Iterable[tuple[int, int, _T]],
/,
) -> None: ...
def filter_edges(self, filter_function: Callable[[_T], bool]) -> EdgeIndices: ...
Expand Down Expand Up @@ -1238,9 +1238,9 @@ class PyGraph(Generic[_S, _T]):
) -> PyGraph: ...
def remove_edge(self, node_a: int, node_b: int, /) -> None: ...
def remove_edge_from_index(self, edge: int, /) -> None: ...
def remove_edges_from(self, index_list: Sequence[tuple[int, int]], /) -> None: ...
def remove_edges_from(self, index_list: Iterable[tuple[int, int]], /) -> None: ...
def remove_node(self, node: int, /) -> None: ...
def remove_nodes_from(self, index_list: Sequence[int], /) -> None: ...
def remove_nodes_from(self, index_list: Iterable[int], /) -> None: ...
def subgraph(self, nodes: Sequence[int], /, preserve_attrs: bool = ...) -> PyGraph[_S, _T]: ...
def substitute_node_with_subgraph(
self,
Expand Down Expand Up @@ -1304,14 +1304,14 @@ class PyDiGraph(Generic[_S, _T]):
def add_edge(self, parent: int, child: int, edge: _T, /) -> int: ...
def add_edges_from(
self,
obj_list: Sequence[tuple[int, int, _T]],
obj_list: Iterable[tuple[int, int, _T]],
/,
) -> list[int]: ...
def add_edges_from_no_data(
self: PyDiGraph[_S, _T | None], obj_list: Sequence[tuple[int, int]], /
self: PyDiGraph[_S, _T | None], obj_list: Iterable[tuple[int, int]], /
) -> list[int]: ...
def add_node(self, obj: _S, /) -> int: ...
def add_nodes_from(self, obj_list: Sequence[_S], /) -> NodeIndices: ...
def add_nodes_from(self, obj_list: Iterable[_S], /) -> NodeIndices: ...
def add_parent(self, child: int, obj: _S, edge: _T, /) -> int: ...
def adj(self, node: int, /) -> dict[int, _T]: ...
def adj_direction(self, node: int, direction: bool, /) -> dict[int, _T]: ...
Expand Down Expand Up @@ -1341,11 +1341,11 @@ class PyDiGraph(Generic[_S, _T]):
def edges(self) -> list[_T]: ...
def edge_subgraph(self, edge_list: Sequence[tuple[int, int]], /) -> PyDiGraph[_S, _T]: ...
def extend_from_edge_list(
self: PyDiGraph[_S | None, _T | None], edge_list: Sequence[tuple[int, int]], /
self: PyDiGraph[_S | None, _T | None], edge_list: Iterable[tuple[int, int]], /
) -> None: ...
def extend_from_weighted_edge_list(
self: PyDiGraph[_S | None, _T],
edge_list: Sequence[tuple[int, int, _T]],
edge_list: Iterable[tuple[int, int, _T]],
/,
) -> None: ...
def filter_edges(self, filter_function: Callable[[_T], bool]) -> EdgeIndices: ...
Expand Down Expand Up @@ -1413,7 +1413,7 @@ class PyDiGraph(Generic[_S, _T]):
) -> PyDiGraph: ...
def remove_edge(self, parent: int, child: int, /) -> None: ...
def remove_edge_from_index(self, edge: int, /) -> None: ...
def remove_edges_from(self, index_list: Sequence[tuple[int, int]], /) -> None: ...
def remove_edges_from(self, index_list: Iterable[tuple[int, int]], /) -> None: ...
def remove_node(self, node: int, /) -> None: ...
def remove_node_retain_edges(
self,
Expand All @@ -1431,7 +1431,7 @@ class PyDiGraph(Generic[_S, _T]):
*,
use_outgoing: bool = ...,
) -> None: ...
def remove_nodes_from(self, index_list: Sequence[int], /) -> None: ...
def remove_nodes_from(self, index_list: Iterable[int], /) -> None: ...
def subgraph(
self, nodes: Sequence[int], /, preserve_attrs: bool = ...
) -> PyDiGraph[_S, _T]: ...
Expand Down
70 changes: 36 additions & 34 deletions src/digraph.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1273,20 +1273,18 @@ impl PyDiGraph {

/// Add new edges to the dag.
///
/// :param list obj_list: A list of tuples of the form
/// :param iterable obj_list: An iterable of tuples of the form
/// ``(parent, child, obj)`` to attach to the graph. ``parent`` and
/// ``child`` are integer indices describing where an edge should be
/// added, and obj is the python object for the edge data.
///
/// :returns: A list of int indices of the newly created edges
/// :rtype: list
#[pyo3(text_signature = "(self, obj_list, /)")]
pub fn add_edges_from(
&mut self,
obj_list: Vec<(usize, usize, PyObject)>,
) -> PyResult<Vec<usize>> {
let mut out_list: Vec<usize> = Vec::with_capacity(obj_list.len());
for obj in obj_list {
pub fn add_edges_from(&mut self, obj_list: Bound<'_, PyAny>) -> PyResult<Vec<usize>> {
let mut out_list = Vec::new();
for py_obj in obj_list.iter()? {
let obj = py_obj?.extract::<(usize, usize, PyObject)>()?;
let edge = self.add_edge(obj.0, obj.1, obj.2)?;
out_list.push(edge);
}
Expand All @@ -1295,7 +1293,7 @@ impl PyDiGraph {

/// Add new edges to the dag without python data.
///
/// :param list obj_list: A list of tuples of the form
/// :param iterable obj_list: An iterable of tuples of the form
/// ``(parent, child)`` to attach to the graph. ``parent`` and
/// ``child`` are integer indices describing where an edge should be
/// added. Unlike :meth:`add_edges_from` there is no data payload and
Expand All @@ -1307,10 +1305,11 @@ impl PyDiGraph {
pub fn add_edges_from_no_data(
&mut self,
py: Python,
obj_list: Vec<(usize, usize)>,
obj_list: Bound<'_, PyAny>,
) -> PyResult<Vec<usize>> {
let mut out_list: Vec<usize> = Vec::with_capacity(obj_list.len());
for obj in obj_list {
let mut out_list = Vec::new();
for py_obj in obj_list.iter()? {
let obj = py_obj?.extract::<(usize, usize)>()?;
let edge = self.add_edge(obj.0, obj.1, py.None())?;
out_list.push(edge);
}
Expand All @@ -1322,17 +1321,18 @@ impl PyDiGraph {
/// This method differs from :meth:`add_edges_from_no_data` in that it will
/// add nodes if a node index is not present in the edge list.
///
/// :param list edge_list: A list of tuples of the form ``(source, target)``
/// :param iterable edge_list: An iterable of tuples of the form ``(source, target)``
/// where source and target are integer node indices. If the node index
/// is not present in the graph, nodes will be added (with a node
/// weight of ``None``) to that index.
#[pyo3(text_signature = "(self, edge_list, /)")]
pub fn extend_from_edge_list(
&mut self,
py: Python,
edge_list: Vec<(usize, usize)>,
edge_list: Bound<'_, PyAny>,
) -> PyResult<()> {
for (source, target) in edge_list {
for py_obj in edge_list.iter()? {
let (source, target) = py_obj?.extract::<(usize, usize)>()?;
let max_index = cmp::max(source, target);
while max_index >= self.node_count() {
self.graph.add_node(py.None());
Expand All @@ -1347,17 +1347,18 @@ impl PyDiGraph {
/// This method differs from :meth:`add_edges_from` in that it will
/// add nodes if a node index is not present in the edge list.
///
/// :param list edge_list: A list of tuples of the form
/// :param iterable edge_list: An iterable of tuples of the form
/// ``(source, target, weight)`` where source and target are integer
/// node indices. If the node index is not present in the graph
/// nodes will be added (with a node weight of ``None``) to that index.
#[pyo3(text_signature = "(self, edge_list, /)")]
pub fn extend_from_weighted_edge_list(
&mut self,
py: Python,
edge_list: Vec<(usize, usize, PyObject)>,
edge_list: Bound<'_, PyAny>,
) -> PyResult<()> {
for (source, target, weight) in edge_list {
for py_obj in edge_list.iter()? {
let (source, target, weight) = py_obj?.extract::<(usize, usize, PyObject)>()?;
let max_index = cmp::max(source, target);
while max_index >= self.node_count() {
self.graph.add_node(py.None());
Expand Down Expand Up @@ -1500,17 +1501,16 @@ impl PyDiGraph {
/// Note if there are multiple edges between the specified nodes only one
/// will be removed.
///
/// :param list index_list: A list of node index pairs to remove from
/// :param iterable index_list: An iterable of node index pairs to remove from
/// the graph
///
/// :raises NoEdgeBetweenNodes: If there are no edges between a specified
/// pair of nodes.
#[pyo3(text_signature = "(self, index_list, /)")]
pub fn remove_edges_from(&mut self, index_list: Vec<(usize, usize)>) -> PyResult<()> {
for (p_index, c_index) in index_list
.iter()
.map(|(x, y)| (NodeIndex::new(*x), NodeIndex::new(*y)))
{
pub fn remove_edges_from(&mut self, index_list: Bound<'_, PyAny>) -> PyResult<()> {
for py_obj in index_list.iter()? {
let (x, y) = py_obj?.extract::<(usize, usize)>()?;
let (p_index, c_index) = (NodeIndex::new(x), NodeIndex::new(y));
let edge_index = match self.graph.find_edge(p_index, c_index) {
Some(edge_index) => edge_index,
None => return Err(NoEdgeBetweenNodes::new_err("No edge found between nodes")),
Expand Down Expand Up @@ -1949,30 +1949,32 @@ impl PyDiGraph {

/// Add new nodes to the graph.
///
/// :param list obj_list: A list of python objects to attach to the graph
/// :param iterable obj_list: An iterable of python objects to attach to the graph
/// as new nodes
///
/// :returns: A list of int indices of the newly created nodes
/// :rtype: NodeIndices
#[pyo3(text_signature = "(self, obj_list, /)")]
pub fn add_nodes_from(&mut self, obj_list: Vec<PyObject>) -> NodeIndices {
let out_list: Vec<usize> = obj_list
.into_iter()
.map(|obj| self.graph.add_node(obj).index())
.collect();
NodeIndices { nodes: out_list }
pub fn add_nodes_from(&mut self, obj_list: Bound<'_, PyAny>) -> PyResult<NodeIndices> {
let mut out_list = Vec::new();
for py_obj in obj_list.iter()? {
let obj = py_obj?.extract::<PyObject>()?;
out_list.push(self.graph.add_node(obj).index());
}
Ok(NodeIndices { nodes: out_list })
}

/// Remove nodes from the graph.
///
/// If a node index in the list is not present in the graph it will be
/// ignored.
///
/// :param list index_list: A list of node indicies to remove from the
/// the graph.
/// :param iterable index_list: An iterable of node indices to remove from the
/// graph.
#[pyo3(text_signature = "(self, index_list, /)")]
pub fn remove_nodes_from(&mut self, index_list: Vec<usize>) -> PyResult<()> {
for node in index_list {
pub fn remove_nodes_from(&mut self, index_list: Bound<'_, PyAny>) -> PyResult<()> {
for py_obj in index_list.iter()? {
let node = py_obj?.extract::<usize>()?;
self.remove_node(node)?;
}
Ok(())
Expand Down
Loading