1 files changed, 19 insertions, 16 deletions
diff --git a/synapse/util/wheel_timer.py b/synapse/util/wheel_timer.py
index 61814aff24..e108adc460 100644
--- a/synapse/util/wheel_timer.py
+++ b/synapse/util/wheel_timer.py
@@ -11,38 +11,41 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Generic, List, TypeVar
 
+T = TypeVar("T")
 
-class _Entry:
+
+class _Entry(Generic[T]):
     __slots__ = ["end_key", "queue"]
 
-    def __init__(self, end_key):
-        self.end_key = end_key
-        self.queue = []
+    def __init__(self, end_key: int) -> None:
+        self.end_key: int = end_key
+        self.queue: List[T] = []
 
 
-class WheelTimer:
+class WheelTimer(Generic[T]):
     """Stores arbitrary objects that will be returned after their timers have
     expired.
     """
 
-    def __init__(self, bucket_size=5000):
+    def __init__(self, bucket_size: int = 5000) -> None:
         """
         Args:
-            bucket_size (int): Size of buckets in ms. Corresponds roughly to the
+            bucket_size: Size of buckets in ms. Corresponds roughly to the
                 accuracy of the timer.
         """
-        self.bucket_size = bucket_size
-        self.entries = []
-        self.current_tick = 0
+        self.bucket_size: int = bucket_size
+        self.entries: List[_Entry[T]] = []
+        self.current_tick: int = 0
 
-    def insert(self, now, obj, then):
+    def insert(self, now: int, obj: T, then: int) -> None:
         """Inserts object into timer.
 
         Args:
-            now (int): Current time in msec
-            obj (object): Object to be inserted
-            then (int): When to return the object strictly after.
+            now: Current time in msec
+            obj: Object to be inserted
+            then: When to return the object strictly after.
         """
         then_key = int(then / self.bucket_size) + 1
 
@@ -70,7 +73,7 @@ class WheelTimer:
 
         self.entries[-1].queue.append(obj)
 
-    def fetch(self, now):
+    def fetch(self, now: int) -> List[T]:
         """Fetch any objects that have timed out
 
         Args:
@@ -87,5 +90,5 @@ class WheelTimer:
 
         return ret
 
-    def __len__(self):
+    def __len__(self) -> int:
         return sum(len(entry.queue) for entry in self.entries)