add type annotations

Author: graingert

MANIFEST.in

@@ -1,2 +1,2 @@
-include README.rst LICENSE CONTRIBUTORS.rst HISTORY.rst
+include README.rst LICENSE CONTRIBUTORS.rst HISTORY.rst src/priority/py.typed
 

docs/source/conf.py

@@ -15,6 +15,8 @@
 import sys
 import os
 
+from typing import List, Dict
+
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
@@ -71,7 +73,7 @@
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
-exclude_patterns = []
+exclude_patterns = []  # type: List[str]
 
 # The reST default role (used for this markup: `text`) to use for all
 # documents.
@@ -194,7 +196,7 @@
 
 # Additional stuff for the LaTeX preamble.
 #'preamble': '',
-}
+}  # type: Dict[str, str]
 
 # Grouping the document tree into LaTeX files. List of tuples
 # (source start file, target name, title,

mypy.ini

@@ -0,0 +1,12 @@
+[mypy]
+strict=true
+warn_unused_configs=true
+show_error_codes=true
+
+[mypy-test_priority]
+allow_untyped_defs=true
+check_untyped_defs=false
+
+[mypy-conftest]
+allow_untyped_defs=true
+check_untyped_defs=false

setup.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-from setuptools import setup, find_packages
+from setuptools import setup, find_packages  # type: ignore[import]
 
 setup(
     name='priority',
@@ -30,4 +30,8 @@
         'Programming Language :: Python :: Implementation :: CPython',
         'Programming Language :: Python :: Implementation :: PyPy',
     ],
+    extras_require={
+        "mypy":  ["mypy>=0.812"],
+    },
+    install_requires=["typing >= 3.7.4.3; python_version < '3.5'"],
 )

src/priority/priority.py

@@ -10,6 +10,7 @@
 import heapq
 import sys
 
+from typing import List, Tuple, Optional
 
 PY3 = sys.version_info[0] == 3
 
@@ -84,22 +85,22 @@ class Stream(object):
     :param stream_id: The stream ID for the new stream.
     :param weight: (optional) The stream weight. Defaults to 16.
     """
-    def __init__(self, stream_id, weight=16):
+    def __init__(self, stream_id, weight=16):  # type: (int, int) -> None
         self.stream_id = stream_id
         self.weight = weight
-        self.children = []
-        self.parent = None
-        self.child_queue = []
+        self.children = []  # type: List[Stream]
+        self.parent = None  # type: Optional[Stream]
+        self.child_queue = []  # type: List[Tuple[int, Stream]]
         self.active = True
         self.last_weight = 0
         self._deficit = 0
 
     @property
-    def weight(self):
+    def weight(self):  # type: () -> int
         return self._weight
 
     @weight.setter
-    def weight(self, value):
+    def weight(self, value):  # type: (int) -> None
         # RFC 7540 § 5.3.2: "All dependent streams are allocated an integer
         # weight between 1 and 256 (inclusive)."
         if not isinstance(value, int):
@@ -109,7 +110,7 @@ def weight(self, value):
                 "Stream weight must be between 1 and 256 (inclusive)")
         self._weight = value
 
-    def add_child(self, child):
+    def add_child(self, child):  # type: (Stream) -> None
         """
         Add a stream that depends on this one.
 
@@ -119,7 +120,7 @@ def add_child(self, child):
         self.children.append(child)
         heapq.heappush(self.child_queue, (self.last_weight, child))
 
-    def add_child_exclusive(self, child):
+    def add_child_exclusive(self, child):  # type: (Stream) -> None
         """
         Add a stream that exclusively depends on this one.
 
@@ -134,7 +135,12 @@ def add_child_exclusive(self, child):
         for old_child in old_children:
             child.add_child(old_child)
 
-    def remove_child(self, child, strip_children=True):
+    def remove_child(
+        self,
+        child,  # type: Stream
+        strip_children=True,  # type: bool
+    ):
+        # type: (...) -> None
         """
         Removes a child stream from this stream. This is a potentially somewhat
         expensive operation.
@@ -150,7 +156,7 @@ def remove_child(self, child, strip_children=True):
         #   it in the old one
         self.children.remove(child)
 
-        new_queue = []
+        new_queue = []  # type: List[Tuple[int, Stream]]
 
         while self.child_queue:
             level, stream = heapq.heappop(self.child_queue)
@@ -165,7 +171,7 @@ def remove_child(self, child, strip_children=True):
             for new_child in child.children:
                 self.add_child(new_child)
 
-    def schedule(self):
+    def schedule(self):  # type: () -> int
         """
         Returns the stream ID of the next child to schedule. Potentially
         recurses down the tree of priorities.
@@ -205,45 +211,45 @@ def schedule(self):
         return next_stream
 
     # Custom repr
-    def __repr__(self):
+    def __repr__(self):  # type: () -> str
         return "Stream<id=%d, weight=%d>" % (self.stream_id, self.weight)
 
     # Custom comparison
-    def __eq__(self, other):
+    def __eq__(self, other):  # type: (object) -> bool
         if not isinstance(other, Stream):  # pragma: no cover
             return False
 
         return self.stream_id == other.stream_id
 
-    def __ne__(self, other):
+    def __ne__(self, other):  # type: (object) -> bool
         return not self.__eq__(other)
 
-    def __lt__(self, other):
+    def __lt__(self, other):  # type: (Stream) -> bool
         if not isinstance(other, Stream):  # pragma: no cover
             return NotImplemented
 
         return self.stream_id < other.stream_id
 
-    def __le__(self, other):
+    def __le__(self, other):  # type: (Stream) -> bool
         if not isinstance(other, Stream):  # pragma: no cover
             return NotImplemented
 
         return self.stream_id <= other.stream_id
 
-    def __gt__(self, other):
+    def __gt__(self, other):  # type: (Stream) -> bool
         if not isinstance(other, Stream):  # pragma: no cover
             return NotImplemented
 
         return self.stream_id > other.stream_id
 
-    def __ge__(self, other):
+    def __ge__(self, other):  # type: (Stream) -> bool
         if not isinstance(other, Stream):  # pragma: no cover
             return NotImplemented
 
         return self.stream_id >= other.stream_id
 
 
-def _stream_cycle(new_parent, current):
+def _stream_cycle(new_parent, current):  # type: (Stream, Stream) -> bool
     """
     Reports whether the new parent depends on the current stream.
     """
@@ -253,7 +259,7 @@ def _stream_cycle(new_parent, current):
     # get more than 100 streams deep. This should catch accidental
     # tree loops. This is the definition of defensive programming.
     for _ in range(100):
-        parent = parent.parent
+        parent = parent.parent  # type: ignore[assignment]
         if parent.stream_id == current.stream_id:
             return True
         elif parent.stream_id == 0:
@@ -289,7 +295,7 @@ class PriorityTree(object):
         default.
     :type maximum_streams: ``int``
     """
-    def __init__(self, maximum_streams=1000):
+    def __init__(self, maximum_streams=1000): # type: (int) -> None
         # This flat array keeps hold of all the streams that are logically
         # dependent on stream 0.
         self._root_stream = Stream(stream_id=0, weight=1)
@@ -302,7 +308,7 @@ def __init__(self, maximum_streams=1000):
             raise ValueError("maximum_streams must be a positive integer.")
         self._maximum_streams = maximum_streams
 
-    def _get_or_insert_parent(self, parent_stream_id):
+    def _get_or_insert_parent(self, parent_stream_id): # type: (int) -> Stream
         """
         When inserting or reprioritizing a stream it is possible to make it
         dependent on a stream that is no longer in the tree. In this situation,
@@ -316,18 +322,26 @@ def _get_or_insert_parent(self, parent_stream_id):
             self.block(parent_stream_id)
             return self._streams[parent_stream_id]
 
-    def _exclusive_insert(self, parent_stream, inserted_stream):
+    def _exclusive_insert(
+        self,
+        parent_stream,  # type: Stream
+        inserted_stream,  # type: Stream
+    ):
+        # type: (...) -> None
         """
         Insert ``inserted_stream`` beneath ``parent_stream``, obeying the
         semantics of exclusive insertion.
         """
         parent_stream.add_child_exclusive(inserted_stream)
 
-    def insert_stream(self,
-                      stream_id,
-                      depends_on=None,
-                      weight=16,
-                      exclusive=False):
+    def insert_stream(
+        self,
+        stream_id,  # type: int
+        depends_on=None,  # type: Optional[int]
+        weight=16,  # type: int
+        exclusive=False, # type: bool
+    ):
+        # type: (...) -> None
         """
         Insert a stream into the tree.
 
@@ -368,11 +382,14 @@ def insert_stream(self,
         parent.add_child(stream)
         self._streams[stream_id] = stream
 
-    def reprioritize(self,
-                     stream_id,
-                     depends_on=None,
-                     weight=16,
-                     exclusive=False):
+    def reprioritize(
+        self,
+        stream_id,  # type: int
+        depends_on=None,  # type: Optional[int]
+        weight=16,  # type: int
+        exclusive=False,  # type: bool
+    ):
+        # type: (...) -> None
         """
         Update the priority status of a stream already in the tree.
 
@@ -415,10 +432,10 @@ def reprioritize(self,
         # its parent, and make it a child of our current parent, and then
         # continue.
         if cycle:
-            new_parent.parent.remove_child(new_parent)
-            current_stream.parent.add_child(new_parent)
+            new_parent.parent.remove_child(new_parent)  # type: ignore[union-attr]
+            current_stream.parent.add_child(new_parent)  # type: ignore[union-attr]
 
-        current_stream.parent.remove_child(
+        current_stream.parent.remove_child(  # type: ignore[union-attr]
             current_stream, strip_children=False
         )
 
@@ -427,7 +444,7 @@ def reprioritize(self,
         else:
             new_parent.add_child(current_stream)
 
-    def remove_stream(self, stream_id):
+    def remove_stream(self, stream_id):  # type: (int) -> None
         """
         Removes a stream from the priority tree.
 
@@ -442,9 +459,10 @@ def remove_stream(self, stream_id):
             raise MissingStreamError("Stream %d not in tree" % stream_id)
 
         parent = child.parent
-        parent.remove_child(child)
+        parent.remove_child(child)  # type: ignore[union-attr]
+
 
-    def block(self, stream_id):
+    def block(self, stream_id): # type: (int) -> None
         """
         Marks a given stream as blocked, with no data to send.
 
@@ -458,7 +476,7 @@ def block(self, stream_id):
         except KeyError:
             raise MissingStreamError("Stream %d not in tree" % stream_id)
 
-    def unblock(self, stream_id):
+    def unblock(self, stream_id):  # type: (int) -> None
         """
         Marks a given stream as unblocked, with more data to send.
 
@@ -473,14 +491,14 @@ def unblock(self, stream_id):
             raise MissingStreamError("Stream %d not in tree" % stream_id)
 
     # The iterator protocol
-    def __iter__(self):  # pragma: no cover
+    def __iter__(self):  # type: () -> PriorityTree  # pragma: no cover
         return self
 
-    def __next__(self):  # pragma: no cover
+    def __next__(self):  # type: () -> int  # pragma: no cover
         try:
             return self._root_stream.schedule()
         except IndexError:
             raise DeadlockError("No unblocked streams to schedule.")
 
-    def next(self):  # pragma: no cover
+    def next(self):  # type: () -> int  # pragma: no cover
         return self.__next__()

src/priority/py.typed

@@ -5,7 +5,7 @@
 
 Py.test fixtures for testing Priority.
 """
-import pytest
+import pytest # type: ignore[import]
 
 import priority