Skip to content

Add __class_getitem__ to generic classes #1345

@barakatzir

Description

@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: ...
    ...

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions