Add __class_getitem__ to generic classes

[barakatzir]

I noticed that despite the PyGraph, PyDiGraph and PyDAG classes being typed as generic over node payload and edge payload type the classes cannot be typed in runtime because the do not implement the __class_getitem__ class method.

This means that using them in annotations requires one of the following: (1) use forward reference (i.e. put them quotation marks); (2) be enclosed in an if TYPE_CHECKING: block; (3) have from __future__ import annotations used in the module; (4) used only in .pyi files. (5) use the type statement in python>=3.12 instead of TypeAlias (which does lazy evaluation).

These are a bit of a hassle but not too troubling, but this has another implication: these type annotations cannot be inspected at runtime.
Runtime inspection is something that packages like pydantic, attrs do (even the standard library dataclasses module does a bit of it).

I suggest adding a __class_getitem__ classmethod for these classes. In python its easy to do (although I'm not suggesting implementing this in python):

from types import GenericAlias

class PyGraph:
    @classmethod
    def __class_getitem__(cls, key, /):
        return GenericAlias(cls, key)
    ...

This should probably be implemented in rust via PyO3. The GenericAlias class has a python-stable ABI counterpart Py_GenericAlias and I assume that it's supported by PyO3 (but haven't checked).

For reference you can look at numpy C-implementation, where they do minimal check that the key is either a tuple of length 1 or 2 or any other python object (the Py_GenericAlias converts it to a one-tuple in that case): https://github.com/numpy/numpy/blob/ebbd9442e9ccd2508979d58354d16e3eeca96d9c/numpy/_core/src/multiarray/methods.c#L2805

If you do add the __class_getitem__ method, then its probably also best to add its signature to the pyi files:

class PyGraph(Generic[_S, _T]):
    def __class_getitem__(cls, item: Any, /) -> GenericAlias: ...
    ...

[IvanIsCoding]

Thanks for opening this. I will give some context that:

• Our code is manually annotated with .pyi files
• PyO3 classes cannot inherit from Python classes

So rustworkx tries to mimic a Python implementations the most possible despite being written in Rust. But some more advanced features are missing.

I agree that this should be available and I will try to implement it. I will also accept PRs if someone wants to do it before I do. The NumPy code is a good starting point.

In the meantime, we will keep suggesting users to stick with temporary fix n. 3 that is from __future__ import annotations

[IvanIsCoding]

#1348 will add the support for PyGraph and PyDiGraph which will cover most use cases

I still need to figure out how to implement this for the Custom Return Types as not all of them can be parametrized

[barakatzir]

Very cool and fast handling! I skimmed the PR. I was kind of let down by the fact the PyO3 doesn't support Py_GenericAlias, so you need to import GenericAlias from types. Do notice that this means that a user will be able to call:

graph = PyGraph[int, str]()

which will need to import types and might cause slowdown on first such import. If somebody asks, the following should be faster

from __future__ import annotations
graph: PyGraph[int, str] = PyGraph()

But I don't think its too bad either way.

[IvanIsCoding]

I need to open an issue in PyO3 asking for the Py_GenericAlias binding to Rust.

It would save us one Python import, which can have a positive performance impact as you mentioned.

However at the moment I think it is a reasonable fix as it is not a breaking change. Code without __future__ already errors anyway