Merge branch 'more-typing' into 'master'

Add more typing

See merge request poezio/slixmpp!166

1 files changed, 31 insertions, 14 deletions

diff --git a/slixmpp/xmlstream/handler/waiter.py b/slixmpp/xmlstream/handler/waiter.py
index 758cd1f1..dde49754 100644
--- a/slixmpp/xmlstream/handler/waiter.py
+++ b/slixmpp/xmlstream/handler/waiter.py
@@ -4,13 +4,20 @@
 # Part of Slixmpp: The Slick XMPP Library
 # :copyright: (c) 2011 Nathanael C. Fritz
 # :license: MIT, see LICENSE for more details
+from __future__ import annotations
+
 import logging
-import asyncio
-from  asyncio import Queue, wait_for, TimeoutError
+from asyncio import Event, wait_for, TimeoutError
+from typing import Optional, TYPE_CHECKING, Union
+from xml.etree.ElementTree import Element
 
 import slixmpp
+from slixmpp.xmlstream.stanzabase import StanzaBase
 from slixmpp.xmlstream.handler.base import BaseHandler
+from slixmpp.xmlstream.matcher.base import MatcherBase
 
+if TYPE_CHECKING:
+    from slixmpp.xmlstream.xmlstream import XMLStream
 
 log = logging.getLogger(__name__)
 
@@ -28,24 +35,27 @@ class Waiter(BaseHandler):
     :param stream: The :class:`~slixmpp.xmlstream.xmlstream.XMLStream`
                    instance this handler should monitor.
     """
+    _event: Event
 
-    def __init__(self, name, matcher, stream=None):
+    def __init__(self, name: str, matcher: MatcherBase, stream: Optional[XMLStream] = None):
         BaseHandler.__init__(self, name, matcher, stream=stream)
-        self._payload = Queue()
+        self._event = Event()
 
-    def prerun(self, payload):
+    def prerun(self, payload: StanzaBase) -> None:
         """Store the matched stanza when received during processing.
 
         :param payload: The matched
-            :class:`~slixmpp.xmlstream.stanzabase.ElementBase` object.
+            :class:`~slixmpp.xmlstream.stanzabase.StanzaBase` object.
         """
-        self._payload.put_nowait(payload)
+        if not self._event.is_set():
+            self._event.set()
+            self._payload = payload
 
-    def run(self, payload):
+    def run(self, payload: StanzaBase) -> None:
         """Do not process this handler during the main event loop."""
         pass
 
-    async def wait(self, timeout=None):
+    async def wait(self, timeout: Optional[int] = None) -> Optional[StanzaBase]:
         """Block an event handler while waiting for a stanza to arrive.
 
         Be aware that this will impact performance if called from a
@@ -59,17 +69,24 @@ class Waiter(BaseHandler):
             :class:`~slixmpp.xmlstream.xmlstream.XMLStream.response_timeout`
             value.
         """
+        stream_ref = self.stream
+        if stream_ref is None:
+            raise ValueError('wait() called without a stream')
+        stream = stream_ref()
+        if stream is None:
+            raise ValueError('wait() called without a stream')
         if timeout is None:
             timeout = slixmpp.xmlstream.RESPONSE_TIMEOUT
 
-        stanza = None
         try:
-            stanza = await self._payload.get()
+            await wait_for(
+                self._event.wait(), timeout, loop=stream.loop
+            )
         except TimeoutError:
             log.warning("Timed out waiting for %s", self.name)
-        self.stream().remove_handler(self.name)
-        return stanza
+        stream.remove_handler(self.name)
+        return self._payload
 
-    def check_delete(self):
+    def check_delete(self) -> bool:
         """Always remove waiters after use."""
         return True