Add capacity keyword arguments to graph creation

releasenotes/notes/with-capacity-6f35d43e256e268e.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    :class:`.PyGraph` and :class:`.PyDiGraph` now each take keyword arguments ``initial_node_count``
+    and ``initial_edge_count`` in their constructors, which can be used to pre-allocate space for
+    the given number of nodes and edges.

rustworkx/rustworkx.pyi

@@ -1209,6 +1209,7 @@ class PyGraph(Generic[_S, _T]):
     ) -> None: ...
     def __delitem__(self, idx: int, /) -> None: ...
     def __getitem__(self, idx: int, /) -> _S: ...
+    def __getnewargs_ex__(self) -> tuple[tuple[Any, ...], dict[str, Any]]: ...
     def __getstate__(self) -> Any: ...
     def __len__(self) -> int: ...
     def __setitem__(self, idx: int, value: _S, /) -> None: ...
@@ -1225,6 +1226,10 @@ class PyDiGraph(Generic[_S, _T]):
         /,
         check_cycle: bool = ...,
         multigraph: bool = ...,
+        attrs: object = ...,
+        *,
+        initial_node_count: int = ...,
+        initial_edge_count: int = ...,
     ) -> None: ...
     def add_child(self, parent: int, obj: _S, edge: _T, /) -> int: ...
     def add_edge(self, parent: int, child: int, edge: _T, /) -> int: ...
@@ -1404,6 +1409,7 @@ class PyDiGraph(Generic[_S, _T]):
     def reverse(self) -> None: ...
     def __delitem__(self, idx: int, /) -> None: ...
     def __getitem__(self, idx: int, /) -> _S: ...
+    def __getnewargs_ex__(self) -> tuple[tuple[Any, ...], dict[str, Any]]: ...
     def __getstate__(self) -> Any: ...
     def __len__(self) -> int: ...
     def __setitem__(self, idx: int, value: _S, /) -> None: ...

src/digraph.rs

@@ -31,7 +31,7 @@ use smallvec::SmallVec;
 use pyo3::exceptions::PyIndexError;
 use pyo3::gc::PyVisit;
 use pyo3::prelude::*;
-use pyo3::types::{PyBool, PyDict, PyList, PyString, PyTuple};
+use pyo3::types::{IntoPyDict, PyBool, PyDict, PyList, PyString, PyTuple};
 use pyo3::PyTraverseError;
 use pyo3::Python;
 
@@ -174,6 +174,10 @@ use super::dag_algo::is_directed_acyclic_graph;
 /// :param attrs: An optional attributes payload to assign to the
 ///     :attr:`~.PyDiGraph.attrs` attribute. This can be any Python object. If
 ///     it is not specified :attr:`~.PyDiGraph.attrs` will be set to ``None``.
+/// :param int initial_node_count: The graph will be allocated with enough capacity to store this
+///     many nodes before needing to grow (default 0: no overallocation).
+/// :param int initial_edge_count: The graph will be allocated with enough capacity to store this
+///     many edges before needing to grow (default 0: no overallocation).
 #[pyclass(mapping, module = "rustworkx", subclass)]
 #[derive(Clone)]
 pub struct PyDiGraph {
@@ -286,10 +290,17 @@ impl PyDiGraph {
 #[pymethods]
 impl PyDiGraph {
     #[new]
-    #[pyo3(signature=(check_cycle=false, multigraph=true, attrs=None), text_signature="(/, check_cycle=False, multigraph=True, attrs=None)")]
-    fn new(py: Python, check_cycle: bool, multigraph: bool, attrs: Option<PyObject>) -> Self {
+    #[pyo3(signature=(/, check_cycle=false, multigraph=true, attrs=None, *, initial_node_count=0, initial_edge_count=0))]
+    fn new(
+        py: Python,
+        check_cycle: bool,
+        multigraph: bool,
+        attrs: Option<PyObject>,
+        initial_node_count: usize,
+        initial_edge_count: usize,
+    ) -> Self {
         PyDiGraph {
-            graph: StablePyGraph::<Directed>::new(),
+            graph: StablePyGraph::<Directed>::with_capacity(initial_node_count, initial_edge_count),
             cycle_state: algo::DfsSpace::default(),
             check_cycle,
             node_removed: false,
@@ -298,6 +309,17 @@ impl PyDiGraph {
         }
     }
 
+    fn __getnewargs_ex__<'py>(&self, py: Python<'py>) -> (Py<PyTuple>, Bound<'py, PyDict>) {
+        (
+            (self.check_cycle, self.multigraph, self.attrs.clone_ref(py)).into_py(py),
+            [
+                ("initial_node_count", self.graph.node_count()),
+                ("initial_edge_count", self.graph.edge_bound()),
+            ]
+            .into_py_dict_bound(py),
+        )
+    }
+
     fn __getstate__(&self, py: Python) -> PyResult<PyObject> {
         let mut nodes: Vec<PyObject> = Vec::with_capacity(self.graph.node_count());
         let mut edges: Vec<PyObject> = Vec::with_capacity(self.graph.edge_bound());
@@ -326,9 +348,6 @@ impl PyDiGraph {
         out_dict.set_item("nodes", nodes_lst)?;
         out_dict.set_item("edges", edges_lst)?;
         out_dict.set_item("nodes_removed", self.node_removed)?;
-        out_dict.set_item("multigraph", self.multigraph)?;
-        out_dict.set_item("attrs", self.attrs.clone_ref(py))?;
-        out_dict.set_item("check_cycle", self.check_cycle)?;
         Ok(out_dict.into())
     }
 
@@ -340,26 +359,11 @@ impl PyDiGraph {
         let edges_lst = binding.downcast::<PyList>()?;
         self.graph = StablePyGraph::<Directed>::new();
         let dict_state = state.downcast_bound::<PyDict>(py)?;
-        self.multigraph = dict_state
-            .get_item("multigraph")?
-            .unwrap()
-            .downcast::<PyBool>()?
-            .extract()?;
         self.node_removed = dict_state
             .get_item("nodes_removed")?
             .unwrap()
             .downcast::<PyBool>()?
             .extract()?;
-        let attrs = match dict_state.get_item("attrs")? {
-            Some(attr) => attr.into(),
-            None => py.None(),
-        };
-        self.attrs = attrs;
-        self.check_cycle = dict_state
-            .get_item("check_cycle")?
-            .unwrap()
-            .downcast::<PyBool>()?
-            .extract()?;
 
         // graph is empty, stop early
         if nodes_lst.is_empty() {

src/graph.rs

@@ -26,7 +26,7 @@ use rustworkx_core::dictmap::*;
 use pyo3::exceptions::PyIndexError;
 use pyo3::gc::PyVisit;
 use pyo3::prelude::*;
-use pyo3::types::{PyBool, PyDict, PyList, PyString, PyTuple};
+use pyo3::types::{IntoPyDict, PyBool, PyDict, PyList, PyString, PyTuple};
 use pyo3::PyTraverseError;
 use pyo3::Python;
 
@@ -138,6 +138,10 @@ use petgraph::visit::{
 /// :param attrs: An optional attributes payload to assign to the
 ///     :attr:`~.PyGraph.attrs` attribute. This can be any Python object. If
 ///     it is not specified :attr:`~.PyGraph.attrs` will be set to ``None``.
+/// :param int initial_node_count: The graph will be allocated with enough capacity to store this
+///     many nodes before needing to grow (default 0: no overallocation).
+/// :param int initial_edge_count: The graph will be allocated with enough capacity to store this
+///     many edges before needing to grow (default 0: no overallocation).
 #[pyclass(mapping, module = "rustworkx", subclass)]
 #[derive(Clone)]
 pub struct PyGraph {
@@ -183,16 +187,36 @@ impl PyGraph {
 #[pymethods]
 impl PyGraph {
     #[new]
-    #[pyo3(signature=(multigraph=true, attrs=None), text_signature = "(/, multigraph=True, attrs=None)")]
-    fn new(py: Python, multigraph: bool, attrs: Option<PyObject>) -> Self {
+    #[pyo3(signature=(multigraph=true, attrs=None, *, initial_node_count=0, initial_edge_count=0))]
+    fn new(
+        py: Python,
+        multigraph: bool,
+        attrs: Option<PyObject>,
+        initial_node_count: usize,
+        initial_edge_count: usize,
+    ) -> Self {
         PyGraph {
-            graph: StablePyGraph::<Undirected>::default(),
+            graph: StablePyGraph::<Undirected>::with_capacity(
+                initial_node_count,
+                initial_edge_count,
+            ),
             node_removed: false,
             multigraph,
             attrs: attrs.unwrap_or_else(|| py.None()),
         }
     }
 
+    fn __getnewargs_ex__<'py>(&self, py: Python<'py>) -> (Py<PyTuple>, Bound<'py, PyDict>) {
+        (
+            (self.multigraph, self.attrs.clone_ref(py)).into_py(py),
+            [
+                ("initial_node_count", self.graph.node_count()),
+                ("initial_edge_count", self.graph.edge_bound()),
+            ]
+            .into_py_dict_bound(py),
+        )
+    }
+
     fn __getstate__(&self, py: Python) -> PyResult<PyObject> {
         let mut nodes: Vec<PyObject> = Vec::with_capacity(self.graph.node_count());
         let mut edges: Vec<PyObject> = Vec::with_capacity(self.graph.edge_bound());
@@ -222,8 +246,6 @@ impl PyGraph {
         out_dict.set_item("nodes", nodes_lst)?;
         out_dict.set_item("edges", edges_lst)?;
         out_dict.set_item("nodes_removed", self.node_removed)?;
-        out_dict.set_item("multigraph", self.multigraph)?;
-        out_dict.set_item("attrs", self.attrs.clone_ref(py))?;
         Ok(out_dict.into())
     }
 
@@ -234,21 +256,11 @@ impl PyGraph {
         let binding = dict_state.get_item("edges")?.unwrap();
         let edges_lst = binding.downcast::<PyList>()?;
 
-        self.graph = StablePyGraph::<Undirected>::default();
-        self.multigraph = dict_state
-            .get_item("multigraph")?
-            .unwrap()
-            .downcast::<PyBool>()?
-            .extract()?;
         self.node_removed = dict_state
             .get_item("nodes_removed")?
             .unwrap()
             .downcast::<PyBool>()?
             .extract()?;
-        self.attrs = match dict_state.get_item("attrs")? {
-            Some(attr) => attr.into(),
-            None => py.None(),
-        };
         // graph is empty, stop early
         if nodes_lst.is_empty() {
             return Ok(());

tests/digraph/test_pickle.py

@@ -30,7 +30,7 @@ def test_noweight_graph(self):
         self.assertEqual({1: (1, 2, None), 3: (3, 1, None)}, dict(gprime.edge_index_map()))
 
     def test_weight_graph(self):
-        g = rx.PyDAG()
+        g = rx.PyDAG(initial_node_count=4, initial_edge_count=4)
         g.add_nodes_from(["A", "B", "C", "D"])
         g.add_edges_from([(0, 1, "A -> B"), (1, 2, "B -> C"), (3, 0, "D -> A"), (3, 1, "D -> B")])
         g.remove_node(0)

tests/graph/test_pickle.py

@@ -30,7 +30,7 @@ def test_noweight_graph(self):
         self.assertEqual({1: (1, 2, None), 3: (3, 1, None)}, dict(gprime.edge_index_map()))
 
     def test_weight_graph(self):
-        g = rx.PyGraph()
+        g = rx.PyGraph(initial_node_count=4, initial_edge_count=4)
         g.add_nodes_from(["A", "B", "C", "D"])
         g.add_edges_from([(0, 1, "A -> B"), (1, 2, "B -> C"), (3, 0, "D -> A"), (3, 1, "D -> B")])
         g.remove_node(0)