Skip to content

networkx: generic typing for node and edge data #15659

@ThePiep

Description

@ThePiep

Currently there are no proper type hints for the data carried by nodes and edges.

Using generics, I propose we get the following functionality:

class NodeData(TypedDict):
    x: int
    y: int

class EdgeData(TypedDict):
    id: int
    color: str

G = nx.Graph[int, NodeData, EdgeData]()

G.add_nodes_from([(1, NodeData(x=0, y=0)), (2, NodeData(x=25, y=30))])
nodes = G.nodes(data=True)
data = nodes[3] # data: NodeData

G.add_edges_from([(1, 2, EdgeData(id=5, color="red"))])

edges = G.edges(data=True)
data = G.get_edge_data(1, 2) # data: EdgeData


G.add_edges_from([(1, 2, EdgeData(id=5, color="red"))]) # correct
G.add_edges_from([(1, 2, EdgeData(id="5", color="red"))]) # error on id
G.add_edges_from([(1, 2, EdgeData(x__coord=10, y__coord=20))]) # error on coords

The default could be set to dict[str,Any], this should also keep compatibility with existing projects as it allows for nx.Graph[Node]() and matches most existing stubs.

Unfortunately, it does not seem possible to enforce these types on insertion. In the add_node() and add_edge() functions due to **attr in their signatures, Unpack does not work for generics. In add_nodes_from() and add_edges_from() we must allow _Node and _Edge respectively in the ebunch (although we could enforce _NodeWithData and _EdgeWithData to facilitate a fully typed workflow, but that would differ from the implementation).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions