Add annotation stubs for PyGraph and PyDiGraph (#401)

* Initial check-in of stub files

* Add types to PyDiGraph

* Remove Union types

* Start adding types to PyGraph

* Tweak __init__

* Finish adding types for PyGraph

* Move PyGraph and PyDiGraph to own stub files

* Import graph annotations in generators module

* Start adding types to custom return types

* Conclude adding types to custom_return_types

* Move stubs to their own package

* Declare package-dir for retworkx-stubs

* Use data_files in setup.py

* Remove print in setup.py

* Fix minor styling issues

* Add simple stubs test case

* Run stub-tests in CI

* Only run stub checks on linux with Python >= 3.7

* Fix error in workflow

* Try to fix 3.6 build

* Improve custom return types

* Add PyDiGraph test

* Add new merge functions

* Add retworkx-stub to MANIFEST.in

* Move stubs to retworkx folder

* Reduce scope of the PR

* Revert changes in setup.py

* Pin mypy version

* Minor improvements for annotations

* Add release note

* Add section to CONTRIBUTING.md about annotations

* Apply suggestions from code review

Co-authored-by: Jake Lishman <jake@binhbar.com>

* Format with Black

* Update dependencies

* Correct pytest-mypy-testing version

* Move tests to new folder

* Black fmt

* Add reveal_type to builtins to help flake8

* Fix path to tests

* Fix error due to line split

* Do not test stubs on Python 3.7

* Add test cases handling restricted methods

* Handle lint

* Use Sequence instead of lists

* Minor updates in tests

* Completely rewrite custom type annotations

* Make PyDAG be generic

* Shorten custom return types

* Correct extend_from_weighted_edge_list signature

* Fix signature

* Yet another signature fix

* Use Self from typing_extensions

* Address PR comments

* Address lint

* Add __str__ and from_complex_adjacency_matrix

* Rename to rustworkx

* Change name of in custom_return_types.pyi

* Skip tests if Python < 3.8

* Ignore no tests found

* Try using pytest==7.1.3

* Improvements after running mypy.stubtest

* More improvements after mypy.stubtest

* Make .pyi file names identical to Rust names

* Make mypy.stubtest happy

* Use mypy.stubtest on CI

* Run Black

* Remove things that we do not need anymore

* More things I forgot to remove

* Skip Python 3.7 for mypy

* Add test stubs without Python 3.7

* Don't use latest mypy release

---------

Co-authored-by: Jake Lishman <jake@binhbar.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

.github/workflows/main.yml

@@ -90,6 +90,31 @@ jobs:
         if: runner.os == 'Linux'
       - name: 'Run tests'
         run: tox -epy
+  tests_stubs:
+    needs: [tests]
+    name: python-stubs-${{ matrix.python-version }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.8, 3.9, "3.10", "3.11"]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+          architecture: x64
+      - name: Install Rust toolchain
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: stable
+          target: x86_64-unknown-linux-gnu
+          profile: minimal
+          default: true
+      - name: 'Install dependencies'
+        run: python -m pip install --upgrade 'tox<4'
+      - name: 'Run rustworkx stub tests'
+        run: tox -estubs
   tests_retworkx_compat:
     needs: [build_lint]
     name: python${{ matrix.python-version }}-${{ matrix.platform.python-architecture }} ${{ matrix.platform.os }} ${{ matrix.msrv }}

CONTRIBUTING.md

@@ -272,6 +272,21 @@ 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 `rustworkx` 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.
+
+```
+tox -estubs
+```
+
 ### 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 rustworkx-core *
+recursive-include rustworkx *.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>`__.

rustworkx/__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 src/lib.rs
+
+from .rustworkx import *
+from typing import Generic, TypeVar
+
+S = TypeVar("S")
+T = TypeVar("T")
+
+class PyDAG(Generic[S, T], PyDiGraph[S, T]): ...

rustworkx/digraph.pyi

@@ -0,0 +1,174 @@
+# 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 src/digraph.rs
+
+import numpy as np
+from .iterators import *
+from .graph import PyGraph
+
+from typing import Any, Callable, Dict, Generic, TypeVar, Optional, List, Tuple, Sequence
+
+__all__ = ["PyDiGraph"]
+
+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[T]: ...
+    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]: ...
+    @staticmethod
+    def from_adjacency_matrix(
+        matrix: np.ndarray, /, null_value: float = ...
+    ) -> PyDiGraph[int, float]: ...
+    @staticmethod
+    def from_complex_adjacency_matrix(
+        matrix: np.ndarray, /, null_value: complex = ...
+    ) -> PyDiGraph[int, complex]: ...
+    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]: ...
+    @staticmethod
+    def read_edge_list(
+        path: str,
+        /,
+        comment: Optional[str] = ...,
+        deliminator: Optional[str] = ...,
+        labels: bool = ...,
+    ) -> 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], /, preserve_attrs: bool = ...) -> 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: ...

rustworkx/graph.pyi

@@ -0,0 +1,129 @@
+# 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 src/graph.rs
+
+import numpy as np
+from .iterators import *
+
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    TypeVar,
+    Optional,
+    List,
+    Tuple,
+    Sequence,
+)
+
+__all__ = ["PyGraph"]
+
+S = TypeVar("S")
+T = TypeVar("T")
+
+class PyGraph(Generic[S, T]):
+    multigraph: bool = ...
+    def __init__(self, /, multigraph: bool = ...) -> None: ...
+    def add_edge(self, node_a: int, node_b: 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: PyGraph[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 adj(self, node: int, /) -> Dict[int, T]: ...
+    def compose(
+        self,
+        other: PyGraph[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) -> PyGraph[S, T]: ...
+    def degree(self, node: int, /) -> int: ...
+    def edge_index_map(self) -> EdgeIndexMap[T]: ...
+    def edge_indices(self) -> EdgeIndices: ...
+    def edge_list(self) -> EdgeList: ...
+    def edges(self) -> List[T]: ...
+    def extend_from_edge_list(
+        self: PyGraph[Optional[S], Optional[T]], edge_list: Sequence[Tuple[int, int]], /
+    ) -> None: ...
+    def extend_from_weighted_edge_list(
+        self: PyGraph[Optional[S], T],
+        edge_list: Sequence[Tuple[int, int, T]],
+        /,
+    ) -> None: ...
+    @staticmethod
+    def from_adjacency_matrix(
+        matrix: np.ndarray, /, null_value: float = ...
+    ) -> PyGraph[int, float]: ...
+    @staticmethod
+    def from_complex_adjacency_matrix(
+        matrix: np.ndarray, /, null_value: complex = ...
+    ) -> PyGraph[int, complex]: ...
+    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 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: ...
+    @staticmethod
+    def read_edge_list(
+        path: str,
+        /,
+        comment: Optional[str] = ...,
+        deliminator: Optional[str] = ...,
+        labels: bool = ...,
+    ) -> PyGraph: ...
+    def remove_edge(self, node_a: int, node_b: 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_nodes_from(self, index_list: Sequence[int], /) -> None: ...
+    def subgraph(self, nodes: Sequence[int], /, preserve_attrs: bool = ...) -> PyGraph[S, T]: ...
+    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 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: ...