Add annotation stubs for PyGraph and PyDiGraph

.github/workflows/main.yml

@@ -39,7 +39,7 @@ jobs:
       - name: Black Codestyle Format
         run: black --check --diff retworkx tests
       - name: Python Lint
-        run: flake8 --per-file-ignores='retworkx/__init__.py:F405,F403' setup.py retworkx tests
+        run: flake8 --per-file-ignores='retworkx/__init__.py:F405,F403 tests/stubs/*:E501' setup.py retworkx tests
       - name: retworkx-core Rust Tests
         run: pushd retworkx-core && cargo test && popd
       - name: retworkx-core Docs
@@ -88,6 +88,9 @@ jobs:
         if: runner.os == 'Linux'
       - name: 'Run tests'
         run: tox -epy
+      - name: 'Run stubs-tests'
+        run: tox -estubs
+        if: runner.os == 'Linux' && ${{ matrix.python-version }} != 3.7
   coverage:
     needs: [tests]
     name: Coverage

CONTRIBUTING.md

@@ -272,6 +272,25 @@ web browser by running:
 cargo doc --open
 ```
 
+### Type Annotations
+
+If you have added new methods or changed method signatures, we encourage you to add annotations for 
+those methods in stub files. The stub files are in the `retworkx` directory and have a `.pyi` file extension.
+They contain annotated signatures for Python functions, stripped of their implementation.
+
+While this step is optional, it is very helpful for end-users. Adding annotations lets users type check
+their code with [mypy](http://mypy-lang.org/), which can be helpful for finding bugs.
+
+Just like with tests for the code, annotations are also tested via tox. On Linux and Mac OS
+machines, the tests can be executed via:
+
+```
+tox -estubs
+```
+
+We also encourage type annotation contributions
+to add tests in the `tests/stubs` directory to verify the correctness of the annotations.
+
 ### Release Notes
 
 It is important to document any end user facing changes when we release a new

MANIFEST.in

@@ -4,3 +4,4 @@ include Cargo.toml
 include Cargo.lock
 recursive-include src *
 recursive-include retworkx-core *
+recursive-include retworkx *.pyi py.typed

releasenotes/notes/graph-annotations-1d436930bf60c5c2.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Added type annotations to :class:`~retworkx.PyDiGraph` and
+    :class:`~retworkx.PyGraph`. They can now be statically type checked
+    with `mypy <http://mypy-lang.org>`__.

retworkx/__init__.pyi

@@ -0,0 +1,18 @@
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+# This file contains only type annotations for PyO3 functions and classes
+# For implementation details, see __init__.py and lib.rs
+
+from .retworkx import *
+from typing import Generic, TypeVar
+
+S = TypeVar("S")
+T = TypeVar("T")
+
+class PyDAG(Generic[S, T], PyDiGraph[S, T]): ...

retworkx/custom_return_types.pyi

@@ -0,0 +1,60 @@
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+# This file contains only type annotations for PyO3 functions and classes
+# For implementation details, see __init__.py and lib.rs
+
+from typing import Any, Generic, List, Iterable, Mapping, TypeVar, Tuple, overload
+from collections.abc import ABC, Sequence
+from typing_extensions import Self
+
+S = TypeVar("S")
+T_co = TypeVar("T_co", covariant=True)
+
+class RetworkxCustomVecIter(Generic[T_co], Sequence[T_co], ABC):
+    def __init__(self) -> None: ...
+    def __eq__(self, other: Sequence[T_co]) -> bool: ...
+    @overload
+    def __getitem__(self, index: int) -> T_co: ...
+    @overload
+    def __getitem__(self: Self, index: slice) -> Self: ...
+    def __getstate__(self) -> Any: ...
+    def __hash__(self) -> int: ...
+    def __len__(self) -> int: ...
+    def __ne__(self, other: Sequence[T_co]) -> bool: ...
+    def __setstate__(self, state) -> None: ...
+
+class RetworkxCustomHashMapIter(Generic[S, T_co], Mapping[S, T_co], ABC):
+    def __init__(self) -> None: ...
+    def items(self) -> Iterable[Tuple[S, T_co]]: ...
+    def keys(self) -> Iterable[S]: ...
+    def values(self) -> Iterable[T_co]: ...
+    def __contains__(self, other: S) -> bool: ...
+    def __eq__(self, other: Mapping[S, T_co]) -> bool: ...
+    def __getitem__(self, index: S) -> T_co: ...
+    def __getstate__(self) -> Any: ...
+    def __hash__(self) -> int: ...
+    def __iter__(self) -> Iterable[S]: ...
+    def __len__(self) -> int: ...
+    def __ne__(self, other: Mapping[S, T_co]) -> bool: ...
+    def __setstate__(self, state) -> None: ...
+
+class NodeIndices(RetworkxCustomVecIter[int]): ...
+class PathLengthMapping(RetworkxCustomHashMapIter[int, float]): ...
+class PathMapping(RetworkxCustomHashMapIter[int, NodeIndices]): ...
+class AllPairsPathLengthMapping(RetworkxCustomHashMapIter[int, PathLengthMapping]): ...
+class AllPairsPathMapping(RetworkxCustomHashMapIter[int, PathMapping]): ...
+class BFSSuccessors(Generic[T_co], RetworkxCustomVecIter[Tuple[T_co, List[T_co]]]): ...
+class EdgeIndexMap(Generic[T_co], RetworkxCustomHashMapIter[int, Tuple[int, int, T_co]]): ...
+class EdgeIndices(RetworkxCustomVecIter[int]): ...
+class Chains(RetworkxCustomVecIter[EdgeIndices]): ...
+class EdgeList(RetworkxCustomVecIter[Tuple[int, int]]): ...
+class NodeMap(RetworkxCustomHashMapIter[int, int]): ...
+class NodesCountMapping(RetworkxCustomHashMapIter[int, int]): ...
+class Pos2DMapping(RetworkxCustomHashMapIter[int, Tuple[float, float]]): ...
+class WeightedEdgeList(Generic[T_co], RetworkxCustomVecIter[Tuple[int, int, T_co]]): ...

retworkx/py.typed

@@ -0,0 +1 @@
+partial

retworkx/pydigraph.pyi

@@ -0,0 +1,155 @@
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+# This file contains only type annotations for PyO3 functions and classes
+# For implementation details, see __init__.py and lib.rs
+
+import numpy as np
+from .custom_return_types import *
+from .pygraph import PyGraph
+
+from typing import Any, Callable, Dict, Generic, TypeVar, Optional, List, Tuple, Sequence
+
+S = TypeVar("S")
+T = TypeVar("T")
+
+class PyDiGraph(Generic[S, T]):
+    check_cycle: bool = ...
+    multigraph: bool = ...
+    def __init__(
+        self,
+        /,
+        check_cycle: bool = ...,
+        multigraph: bool = ...,
+    ) -> None: ...
+    def add_child(self, parent: int, obj: S, edge: T, /) -> int: ...
+    def add_edge(self, parent: int, child: int, edge: T, /) -> int: ...
+    def add_edges_from(
+        self,
+        obj_list: Sequence[Tuple[int, int, T]],
+        /,
+    ) -> List[int]: ...
+    def add_edges_from_no_data(
+        self: PyDiGraph[S, Optional[T]], obj_list: Sequence[Tuple[int, int]], /
+    ) -> List[int]: ...
+    def add_node(self, obj: S, /) -> int: ...
+    def add_nodes_from(self, obj_list: Sequence[S], /) -> NodeIndices: ...
+    def add_parent(self, child: int, obj: S, edge: T, /) -> int: ...
+    def adj(self, node: int, /) -> Dict[int, T]: ...
+    def adj_direction(self, node: int, direction: bool, /) -> Dict[int, T]: ...
+    def compose(
+        self,
+        other: PyDiGraph[S, T],
+        node_map: Dict[int, Tuple[int, T]],
+        /,
+        node_map_func: Optional[Callable[[S], int]] = ...,
+        edge_map_func: Optional[Callable[[T], int]] = ...,
+    ) -> Dict[int, int]: ...
+    def copy(self) -> PyDiGraph[S, T]: ...
+    def edge_index_map(self) -> EdgeIndexMap: ...
+    def edge_indices(self) -> EdgeIndices: ...
+    def edge_list(self) -> EdgeList: ...
+    def edges(self) -> List[T]: ...
+    def extend_from_edge_list(
+        self: PyDiGraph[Optional[S], Optional[T]], edge_list: Sequence[Tuple[int, int]], /
+    ) -> None: ...
+    def extend_from_weighted_edge_list(
+        self: PyDiGraph[Optional[S], T],
+        edge_list: Sequence[Tuple[int, int, T]],
+        /,
+    ) -> None: ...
+    def find_adjacent_node_by_edge(self, node: int, predicate: Callable[[T], bool], /) -> S: ...
+    def find_node_by_weight(self, obj: Callable[[S], bool]) -> Optional[int]: ...
+    def find_predecessors_by_edge(
+        self, node: int, filter_fn: Callable[[T], bool], /
+    ) -> List[S]: ...
+    def find_successors_by_edge(self, node: int, filter_fn: Callable[[T], bool], /) -> List[S]: ...
+    @classmethod
+    def from_adjacency_matrix(matrix: np.array, /) -> PyDiGraph: ...
+    def get_all_edge_data(self, node_a: int, node_b: int, /) -> List[T]: ...
+    def get_edge_data(self, node_a: int, node_b: int, /) -> T: ...
+    def get_node_data(self, node: int, /) -> S: ...
+    def has_edge(self, node_a: int, node_b: int, /) -> bool: ...
+    def in_degree(self, node: int, /) -> int: ...
+    def in_edges(self, node: int, /) -> WeightedEdgeList[T]: ...
+    def insert_node_on_in_edges(self, node: int, ref_node: int, /) -> None: ...
+    def insert_node_on_in_edges_multiple(self, node: int, ref_nodes: Sequence[int], /) -> None: ...
+    def insert_node_on_out_edges(self, node: int, ref_node: int, /) -> None: ...
+    def insert_node_on_out_edges_multiple(self, node: int, ref_nodes: Sequence[int], /) -> None: ...
+    def is_symmetric(self) -> bool: ...
+    def merge_nodes(self, u: int, v: int, /) -> None: ...
+    def neighbors(self, node: int, /) -> NodeIndices: ...
+    def node_indexes(self) -> NodeIndices: ...
+    def nodes(self) -> List[S]: ...
+    def num_edges(self) -> int: ...
+    def num_nodes(self) -> int: ...
+    def out_degree(self, node: int, /) -> int: ...
+    def out_edges(self, node: int, /) -> WeightedEdgeList[T]: ...
+    def predecessor_indices(self, node: int, /) -> NodeIndices: ...
+    def predecessors(self, node: int, /) -> List[S]: ...
+    @classmethod
+    def read_edge_list(
+        path: str,
+        /,
+        comment: Optional[str] = ...,
+        deliminator: Optional[str] = ...,
+    ) -> PyDiGraph: ...
+    def remove_edge(self, parent: int, child: int, /) -> None: ...
+    def remove_edge_from_index(self, edge: int, /) -> None: ...
+    def remove_edges_from(self, index_list: Sequence[Tuple[int, int]], /) -> None: ...
+    def remove_node(self, node: int, /) -> None: ...
+    def remove_node_retain_edges(
+        self,
+        node: int,
+        /,
+        use_outgoing: Optional[bool] = ...,
+        condition: Optional[Callable[[S, S], bool]] = ...,
+    ) -> None: ...
+    def remove_nodes_from(self, index_list: Sequence[int], /) -> None: ...
+    def subgraph(self, nodes: Sequence[int], /) -> PyDiGraph[S, T]: ...
+    def substitute_node_with_subgraph(
+        self,
+        node: int,
+        other: PyDiGraph[S, T],
+        edge_map_fn: Callable[[int, int, T], Optional[int]],
+        /,
+        node_filter: Optional[Callable[[S], bool]] = ...,
+        edge_weight_map: Optional[Callable[[T], T]] = ...,
+    ) -> NodeMap: ...
+    def successor_indices(self, node: int, /) -> NodeIndices: ...
+    def successors(self, node: int, /) -> List[S]: ...
+    def to_dot(
+        self,
+        /,
+        node_attr: Optional[Callable[[S], Dict[str, str]]] = ...,
+        edge_attr: Optional[Callable[[T], Dict[str, str]]] = ...,
+        graph_attr: Optional[Dict[str, str]] = ...,
+        filename: Optional[str] = ...,
+    ) -> Optional[str]: ...
+    def to_undirected(
+        self,
+        /,
+        multigraph: bool = ...,
+        weight_combo_fn: Optional[Callable[[T, T], T]] = ...,
+    ) -> PyGraph[S, T]: ...
+    def update_edge(self, source: int, target: int, edge: T, /) -> None: ...
+    def update_edge_by_index(self, edge_index: int, edge: T, /) -> None: ...
+    def weighted_edge_list(self) -> WeightedEdgeList[T]: ...
+    def write_edge_list(
+        self,
+        path: str,
+        /,
+        deliminator: Optional[str] = ...,
+        weight_fn: Optional[Callable[[T], str]] = ...,
+    ) -> None: ...
+    def __delitem__(self, idx: int, /) -> None: ...
+    def __getitem__(self, idx: int, /) -> S: ...
+    def __getstate__(self) -> Any: ...
+    def __len__(self) -> int: ...
+    def __setitem__(self, idx: int, value: S, /) -> None: ...
+    def __setstate__(self, state, /) -> None: ...