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)
|
