Use PyGenericAlias directly from PyO3 (#1421)

Author: IvanIsCoding

src/digraph.rs

@@ -32,7 +32,7 @@ use smallvec::SmallVec;
 use pyo3::exceptions::PyIndexError;
 use pyo3::gc::PyVisit;
 use pyo3::prelude::*;
-use pyo3::types::{IntoPyDict, PyBool, PyDict, PyList, PyString, PyTuple, PyType};
+use pyo3::types::{IntoPyDict, PyBool, PyDict, PyGenericAlias, PyList, PyString, PyTuple, PyType};
 use pyo3::IntoPyObjectExt;
 use pyo3::PyTraverseError;
 use pyo3::Python;
@@ -57,8 +57,8 @@ use super::iterators::{
     EdgeIndexMap, EdgeIndices, EdgeList, NodeIndices, NodeMap, WeightedEdgeList,
 };
 use super::{
-    find_node_by_weight, generic_class_getitem, weight_callable, DAGHasCycle, DAGWouldCycle, IsNan,
-    NoEdgeBetweenNodes, NoSuitableNeighbors, NodesRemoved, StablePyGraph,
+    find_node_by_weight, weight_callable, DAGHasCycle, DAGWouldCycle, IsNan, NoEdgeBetweenNodes,
+    NoSuitableNeighbors, NodesRemoved, StablePyGraph,
 };
 
 use super::dag_algo::is_directed_acyclic_graph;
@@ -3440,7 +3440,8 @@ impl PyDiGraph {
         cls: &Bound<'_, PyType>,
         key: &Bound<'_, PyAny>,
     ) -> PyResult<PyObject> {
-        generic_class_getitem(cls, key)
+        let alias = PyGenericAlias::new(cls.py(), cls.as_any(), key)?;
+        Ok(alias.into())
     }
 
     // Functions to enable Python Garbage Collection

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, PyType};
+use pyo3::types::{IntoPyDict, PyBool, PyDict, PyGenericAlias, PyList, PyString, PyTuple, PyType};
 use pyo3::IntoPyObjectExt;
 use pyo3::PyTraverseError;
 use pyo3::Python;
@@ -41,8 +41,7 @@ use crate::iterators::NodeMap;
 use super::dot_utils::build_dot;
 use super::iterators::{EdgeIndexMap, EdgeIndices, EdgeList, NodeIndices, WeightedEdgeList};
 use super::{
-    find_node_by_weight, generic_class_getitem, weight_callable, IsNan, NoEdgeBetweenNodes,
-    NodesRemoved, StablePyGraph,
+    find_node_by_weight, weight_callable, IsNan, NoEdgeBetweenNodes, NodesRemoved, StablePyGraph,
 };
 
 use crate::RxPyResult;
@@ -2058,7 +2057,8 @@ impl PyGraph {
         cls: &Bound<'_, PyType>,
         key: &Bound<'_, PyAny>,
     ) -> PyResult<PyObject> {
-        generic_class_getitem(cls, key)
+        let alias = PyGenericAlias::new(cls.py(), cls.as_any(), key)?;
+        Ok(alias.into())
     }
 
     // Functions to enable Python Garbage Collection

src/lib.rs

@@ -371,19 +371,6 @@ 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("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.