Add stubs for missing packages

.pre-commit-config.yaml

@@ -4,6 +4,10 @@ repos:
     hooks:
       - id: black
         args: ["--line-length", "79"]
+      - id: black
+        name: black (.pyi)
+        args: ["--line-length", "79"]
+        types: [pyi]
   - repo: local
     hooks:
       - id: clang-format

python/CMakeLists.txt

@@ -4,6 +4,7 @@ if(GTIRB_ENABLE_MYPY)
   if(MYPY AND MYPY_PROTOBUF)
     set(ENABLE_MYPY ON)
     set(MYPY_PROTOBUF_FLAGS --mypy_out=${CMAKE_CURRENT_BINARY_DIR})
+    configure_file(mypy.ini.in mypy.ini @ONLY)
   else()
     set(ENABLE_MYPY OFF)
     if(NOT MYPY)

python/gtirb/byteinterval.py

@@ -172,7 +172,7 @@ def __init__(
             raise ValueError("initialized_size must be <= size!")
 
         super().__init__(uuid=uuid)
-        self._interval_tree = IntervalTree()
+        self._interval_tree: "IntervalTree[int, ByteBlock]" = IntervalTree()
         self._section = None  # type: typing.Optional["Section"]
         self.address = address
         self.size = size

python/gtirb/cfg.py

@@ -1,14 +1,22 @@
-import typing
 from enum import Enum
+from typing import (
+    TYPE_CHECKING,
+    Hashable,
+    Iterable,
+    Iterator,
+    MutableSet,
+    NamedTuple,
+    Optional,
+)
 from uuid import UUID
 
-import networkx
+from networkx import MultiDiGraph
 
 from .block import CfgNode
 from .proto import CFG_pb2
 from .util import DeserializationError
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from .ir import IR
 
 
@@ -70,7 +78,7 @@ class Type(Enum):
 
 
 class Label(
-    typing.NamedTuple(
+    NamedTuple(
         "NamedTuple",
         (("type", "Type"), ("conditional", bool), ("direct", bool)),
     )
@@ -115,12 +123,12 @@ def __repr__(self):
 
 
 class Edge(
-    typing.NamedTuple(
+    NamedTuple(
         "NamedTuple",
         (
             ("source", CfgNode),
             ("target", CfgNode),
-            ("label", "typing.Optional[Label]"),
+            ("label", "Optional[Label]"),
         ),
     )
 ):
@@ -135,14 +143,14 @@ class Edge(
     __slots__ = ()
 
     def __new__(cls, source, target, label=None):
-        # type: (CfgNode, CfgNode, typing.Optional[Label]) -> "Edge"
+        # type: (CfgNode, CfgNode, Optional[Label]) -> "Edge"
         return super().__new__(cls, source, target, label)
 
     Type = _Type
     Label = _Label
 
 
-class CFG(typing.MutableSet[Edge]):
+class CFG(MutableSet[Edge]):
     """A control-flow graph for an :class:`IR`. Vertices are
     :class:`CfgNode`\\s, and edges may optionally contain
     :class:`Edge.Label`\\s.
@@ -161,13 +169,15 @@ class CFG(typing.MutableSet[Edge]):
     """
 
     def __init__(self, edges=None):
-        # type: (typing.Iterable[Edge]) -> None
-        self._nxg = networkx.MultiDiGraph()
+        # type: (Iterable[Edge]) -> None
+        self._nxg: "MultiDiGraph[CfgNode, Hashable, Optional[Label]]" = (
+            MultiDiGraph()
+        )
         if edges is not None:
             self.update(edges)
 
     def _edge_key(self, edge):
-        # type: (Edge) -> typing.Optional[typing.Hashable]
+        # type: (Edge) -> Optional[Hashable]
         if edge.source in self._nxg:
             neighbors = self._nxg[edge.source]
             if edge.target in neighbors:
@@ -181,7 +191,7 @@ def __contains__(self, edge):
         return isinstance(edge, Edge) and self._edge_key(edge) is not None
 
     def __iter__(self):
-        # type: () -> typing.Iterator[Edge]
+        # type: () -> Iterator[Edge]
         for s, t, l in self._nxg.edges(data="label"):
             yield Edge(s, t, l)
 
@@ -190,7 +200,7 @@ def __len__(self):
         return len(self._nxg.edges())
 
     def update(self, edges):
-        # type: (typing.Iterable[Edge]) -> None
+        # type: (Iterable[Edge]) -> None
         for edge in edges:
             self.add(edge)
 
@@ -210,20 +220,20 @@ def discard(self, edge):
             self._nxg.remove_edge(edge.source, edge.target, key=key)
 
     def out_edges(self, node):
-        # type: (CfgNode) -> typing.Iterable[Edge]
+        # type: (CfgNode) -> Iterable[Edge]
         if node in self._nxg:
             for s, t, l in self._nxg.out_edges(node, data="label"):
                 yield Edge(s, t, l)
 
     def in_edges(self, node):
-        # type: (CfgNode) -> typing.Iterable[Edge]
+        # type: (CfgNode) -> Iterable[Edge]
         if node in self._nxg:
             for s, t, l in self._nxg.in_edges(node, data="label"):
                 yield Edge(s, t, l)
 
     @classmethod
     def _from_protobuf(cls, edges, ir):
-        # type: (typing.Iterable[CFG_pb2.Edge], typing.Optional[IR]) -> CFG
+        # type: (Iterable[CFG_pb2.Edge], Optional[IR]) -> CFG
         assert ir
 
         def make_edge(ir, edge):
@@ -242,7 +252,7 @@ def make_edge(ir, edge):
                     "CFG: UUID %s is not a CfgNode" % target_uuid
                 )
 
-            label = None  # type: typing.Optional[Label]
+            label = None  # type: Optional[Label]
             if edge.HasField("label"):
                 label = Edge.Label(
                     Edge.Type(edge.label.type),
@@ -255,7 +265,7 @@ def make_edge(ir, edge):
         return CFG(make_edge(ir, edge) for edge in edges)
 
     def _to_protobuf(self):
-        # type: () -> typing.Iterable[CFG_pb2.Edge]
+        # type: () -> Iterable[CFG_pb2.Edge]
         for s, t, l in self._nxg.edges(data="label"):
             proto_edge = CFG_pb2.Edge()
             proto_edge.source_uuid = s.uuid.bytes
@@ -267,7 +277,7 @@ def _to_protobuf(self):
             yield proto_edge
 
     def nx(self):
-        # type: () -> networkx.MultiDiGraph
+        # type: () -> MultiDiGraph
         return self._nxg
 
     def deep_eq(self, other):

python/gtirb/section.py

@@ -100,7 +100,7 @@ def __init__(
         """
 
         super().__init__(uuid)
-        self._interval_index = IntervalTree()
+        self._interval_index: "IntervalTree[int,ByteInterval]" = IntervalTree()
         self._module = None  # type: typing.Optional["Module"]
         self.name = name  # type: str
         self.byte_intervals = Section._ByteIntervalSet(self, byte_intervals)

python/mypy.ini.in

@@ -0,0 +1,2 @@
+[mypy]
+mypy_path = @CMAKE_CURRENT_SOURCE_DIR@/stubs

python/stubs/README.md

@@ -0,0 +1,5 @@
+Stubs for packages that do not include type annotations recognized by mypy.
+
+Generic type stubs were handwritten to enable more precise type checking.
+They are deliberately incomplete, only including the operations used in the
+GTIRB implementation.

python/stubs/intervaltree/__init__.pyi

@@ -0,0 +1,2 @@
+from .interval import Interval as Interval
+from .intervaltree import IntervalTree as IntervalTree

python/stubs/intervaltree/interval.pyi

@@ -0,0 +1,10 @@
+from typing import Generic, TypeVar
+
+PointT = TypeVar("PointT")
+DataT = TypeVar("DataT")
+
+class Interval(Generic[PointT, DataT]):
+    begin: PointT
+    end: PointT
+    data: DataT
+    def __init__(self, begin: PointT, end: PointT, data: DataT): ...

python/stubs/intervaltree/intervaltree.pyi

@@ -0,0 +1,23 @@
+from typing import Iterator, MutableSet, Set, TypeVar, overload
+
+from .interval import Interval
+
+PointT = TypeVar("PointT")
+DataT = TypeVar("DataT")
+
+class IntervalTree(MutableSet[Interval[PointT, DataT]]):
+    def __contains__(self, item: object) -> bool: ...
+    def __iter__(self) -> Iterator[Interval[PointT, DataT]]: ...
+    def __len__(self) -> int: ...
+    def add(self, interval: Interval[PointT, DataT]) -> None: ...
+    def discard(self, interval: Interval[PointT, DataT]) -> None: ...
+    def begin(self) -> PointT: ...
+    def span(self) -> PointT: ...
+    @overload
+    def overlap(
+        self, begin: Interval[PointT, DataT]
+    ) -> Set[Interval[PointT, DataT]]: ...
+    @overload
+    def overlap(
+        self, begin: PointT, end: PointT
+    ) -> Set[Interval[PointT, DataT]]: ...

python/stubs/networkx/__init__.pyi

@@ -0,0 +1 @@
+from networkx.classes import *

python/stubs/networkx/classes/__init__.pyi

@@ -0,0 +1 @@
+from .multidigraph import MultiDiGraph as MultiDiGraph

python/stubs/networkx/classes/multidigraph.pyi

@@ -0,0 +1,63 @@
+from typing import (
+    Collection,
+    Dict,
+    Generic,
+    Iterable,
+    Tuple,
+    TypeVar,
+    overload,
+)
+from typing_extensions import Literal
+
+NodeT = TypeVar("NodeT")
+KeyT = TypeVar("KeyT")
+DataT = TypeVar("DataT")
+
+class MultiDiGraph(Generic[NodeT, KeyT, DataT]):
+    def __contains__(self, n: object) -> bool: ...
+    def __getitem__(
+        self, n: NodeT
+    ) -> Dict[NodeT, Dict[KeyT, Dict[str, DataT]]]: ...
+    def add_edge(
+        self, u: NodeT, v: NodeT, key: KeyT | None = ..., **attr: DataT
+    ) -> KeyT: ...
+    def remove_edge(
+        self, u: NodeT, v: NodeT, key: KeyT | None = ...
+    ) -> None: ...
+    def number_of_edges(
+        self, u: NodeT | None = ..., v: NodeT | None = ...
+    ) -> int: ...
+    def clear(self) -> None: ...
+    # The actual types of edges(), out_edges(), and in_edges() are remarkably
+    # complicated. Depending on various combinations of arguments, they can
+    # retrieve a collection of tuples of two, three, or four elements. These
+    # overloads are specifically just the cases currently used by CFG, and do
+    # not remotely cover all valid combinations of arguments.
+    @overload
+    def edges(
+        self,
+        nbunch: Iterable[NodeT] | NodeT | None = ...,
+        data: Literal[False] = ...,
+    ) -> Collection[Tuple[NodeT, NodeT]]: ...
+    @overload
+    def edges(
+        self,
+        nbunch: Iterable[NodeT] | NodeT | None = ...,
+        *,
+        data: str,
+        default: DataT | None = ...,
+    ) -> Collection[Tuple[NodeT, NodeT, DataT | None]]: ...
+    def out_edges(
+        self,
+        nbunch: Iterable[NodeT] | NodeT | None = ...,
+        *,
+        data: str = ...,
+        default: DataT | None = ...,
+    ) -> Collection[Tuple[NodeT, NodeT, DataT | None]]: ...
+    def in_edges(
+        self,
+        nbunch: Iterable[NodeT] | NodeT | None = ...,
+        *,
+        data: str = ...,
+        default: DataT | None = ...,
+    ) -> Collection[Tuple[NodeT, NodeT, DataT | None]]: ...

python/stubs/sortedcontainers/__init__.pyi

@@ -0,0 +1 @@
+from .sorteddict import SortedDict as SortedDict

python/stubs/sortedcontainers/sorteddict.pyi

@@ -0,0 +1,11 @@
+from typing import MutableMapping, Iterator, TypeVar
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+class SortedDict(MutableMapping[K, V]):
+    def __delitem__(self, key: K) -> None: ...
+    def __getitem__(self, key: K) -> V: ...
+    def __iter__(self) -> Iterator[K]: ...
+    def __len__(self) -> int: ...
+    def __setitem__(self, key: K, value: V) -> None: ...