Allow PyGraph and PyDiGraph to be annotated as generic classes at runtime (#1348)

* Allow PyGraph and PyDiGraph to be annotated as generic classes at runtime

* Black

* Add release notes

* Add __class_getitem__ to stubs

Author: IvanIsCoding

releasenotes/notes/improve-pep560-support-66f059006db8e3b2.yaml

@@ -0,0 +1,21 @@
+---
+features:
+  - |
+    The :class:`PyDiGraph() <rustworkx.PyDiGraph>` and :class:`PyGraph() <rustworkx.PyGraph>`
+    classes now have better support for
+    `PEP 560 <https://peps.python.org/pep-0560/>`__. Building off of the previous
+    releases which introduced type annotations, the following code snippet is now valid:
+
+    .. jupyter-execute::
+
+      import rustworkx as rx
+
+      graph: rx.PyGraph[int, int] = rx.PyGraph()
+
+
+    Previously, users had to rely on post-poned evaluation of type annotations from
+    `PEP 563 <https://peps.python.org/pep-0563/>`__ for annotations to work.
+
+    Refer to `issue 1345 <https://github.com/Qiskit/rustworkx/issues/1345>`__ for
+    more information.
+    

rustworkx/rustworkx.pyi

@@ -10,6 +10,7 @@
 # For implementation details, see __init__.py and src/lib.rs
 
 from .visit import BFSVisitor, DFSVisitor, DijkstraVisitor
+from types import GenericAlias
 from typing import (
     Callable,
     final,
@@ -1308,6 +1309,8 @@ class PyGraph(Generic[_S, _T]):
     ) -> None: ...
     def __delitem__(self, idx: int, /) -> None: ...
     def __getitem__(self, idx: int, /) -> _S: ...
+    @classmethod
+    def __class_getitem__(cls, key: Any, /) -> GenericAlias: ...
     def __getnewargs_ex__(self) -> tuple[tuple[Any, ...], dict[str, Any]]: ...
     def __getstate__(self) -> Any: ...
     def __len__(self) -> int: ...
@@ -1509,6 +1512,8 @@ class PyDiGraph(Generic[_S, _T]):
     def reverse(self) -> None: ...
     def __delitem__(self, idx: int, /) -> None: ...
     def __getitem__(self, idx: int, /) -> _S: ...
+    @classmethod
+    def __class_getitem__(cls, key: Any, /) -> GenericAlias: ...
     def __getnewargs_ex__(self) -> tuple[tuple[Any, ...], dict[str, Any]]: ...
     def __getstate__(self) -> Any: ...
     def __len__(self) -> int: ...

src/digraph.rs

@@ -31,7 +31,7 @@ use smallvec::SmallVec;
 use pyo3::exceptions::PyIndexError;
 use pyo3::gc::PyVisit;
 use pyo3::prelude::*;
-use pyo3::types::{IntoPyDict, PyBool, PyDict, PyList, PyString, PyTuple};
+use pyo3::types::{IntoPyDict, PyBool, PyDict, PyList, PyString, PyTuple, PyType};
 use pyo3::PyTraverseError;
 use pyo3::Python;
 
@@ -55,8 +55,8 @@ use super::iterators::{
     EdgeIndexMap, EdgeIndices, EdgeList, NodeIndices, NodeMap, WeightedEdgeList,
 };
 use super::{
-    find_node_by_weight, weight_callable, DAGHasCycle, DAGWouldCycle, IsNan, NoEdgeBetweenNodes,
-    NoSuitableNeighbors, NodesRemoved, StablePyGraph,
+    find_node_by_weight, generic_class_getitem, weight_callable, DAGHasCycle, DAGWouldCycle, IsNan,
+    NoEdgeBetweenNodes, NoSuitableNeighbors, NodesRemoved, StablePyGraph,
 };
 
 use super::dag_algo::is_directed_acyclic_graph;
@@ -3143,6 +3143,15 @@ impl PyDiGraph {
         }
     }
 
+    #[classmethod]
+    #[pyo3(signature = (key, /))]
+    pub fn __class_getitem__(
+        cls: &Bound<'_, PyType>,
+        key: &Bound<'_, PyAny>,
+    ) -> PyResult<PyObject> {
+        generic_class_getitem(cls, key)
+    }
+
     // Functions to enable Python Garbage Collection
 
     // Function for PyTypeObject.tp_traverse [1][2] used to tell Python what

src/graph.rs

@@ -26,7 +26,7 @@ use rustworkx_core::graph_ext::*;
 use pyo3::exceptions::PyIndexError;
 use pyo3::gc::PyVisit;
 use pyo3::prelude::*;
-use pyo3::types::{IntoPyDict, PyBool, PyDict, PyList, PyString, PyTuple};
+use pyo3::types::{IntoPyDict, PyBool, PyDict, PyList, PyString, PyTuple, PyType};
 use pyo3::PyTraverseError;
 use pyo3::Python;
 
@@ -40,7 +40,8 @@ use crate::iterators::NodeMap;
 use super::dot_utils::build_dot;
 use super::iterators::{EdgeIndexMap, EdgeIndices, EdgeList, NodeIndices, WeightedEdgeList};
 use super::{
-    find_node_by_weight, weight_callable, IsNan, NoEdgeBetweenNodes, NodesRemoved, StablePyGraph,
+    find_node_by_weight, generic_class_getitem, weight_callable, IsNan, NoEdgeBetweenNodes,
+    NodesRemoved, StablePyGraph,
 };
 
 use crate::RxPyResult;
@@ -1994,6 +1995,15 @@ impl PyGraph {
         }
     }
 
+    #[classmethod]
+    #[pyo3(signature = (key, /))]
+    pub fn __class_getitem__(
+        cls: &Bound<'_, PyType>,
+        key: &Bound<'_, PyAny>,
+    ) -> PyResult<PyObject> {
+        generic_class_getitem(cls, key)
+    }
+
     // Functions to enable Python Garbage Collection
 
     // Function for PyTypeObject.tp_traverse [1][2] used to tell Python what

src/lib.rs

@@ -376,6 +376,19 @@ fn find_node_by_weight<Ty: EdgeType>(
     Ok(index)
 }
 
+fn generic_class_getitem(
+    cls: &Bound<'_, pyo3::types::PyType>,
+    key: &Bound<'_, PyAny>,
+) -> PyResult<PyObject> {
+    Python::with_gil(|py| -> PyResult<PyObject> {
+        let types_mod = py.import_bound("types")?;
+        let types_generic_alias = types_mod.getattr("GenericAlias")?;
+        let args = (cls, key);
+        let generic_alias = types_generic_alias.call1(args)?;
+        Ok(generic_alias.into())
+    })
+}
+
 // The provided node is invalid.
 create_exception!(rustworkx, InvalidNode, PyException);
 // Performing this operation would result in trying to add a cycle to a DAG.

tests/test_annotations.py

@@ -0,0 +1,38 @@
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+import rustworkx as rx
+import types
+import unittest
+
+
+class TestAnnotationSubscriptions(unittest.TestCase):
+    def test_digraph(self):
+        graph: rx.PyDiGraph[int, int] = rx.PyDiGraph()
+        self.assertIsInstance(
+            graph.__class_getitem__((int, int)),
+            types.GenericAlias,
+        )
+
+    def test_graph(self):
+        graph: rx.PyGraph[int, int] = rx.PyGraph()
+        self.assertIsInstance(
+            graph.__class_getitem__((int, int)),
+            types.GenericAlias,
+        )
+
+    def test_dag(self):
+        graph: rx.PyDAG[int, int] = rx.PyDAG()
+        self.assertIsInstance(
+            graph.__class_getitem__((int, int)),
+            types.GenericAlias,
+        )