XMLStream: add a forever parameter to process(), defaulting to True, to select whether we want to stop the event loop after a disconnection

3 files changed, 17 insertions, 7 deletions

diff --git a/docs/using_asyncio.rst b/docs/using_asyncio.rst
index c84148e8..55ed7679 100644
--- a/docs/using_asyncio.rst
+++ b/docs/using_asyncio.rst
@@ -53,7 +53,8 @@ Running the event loop
 
 :meth:`.XMLStream.process` is only a thin wrapper on top of
 ``loop.run_forever()`` (if ``timeout`` is provided then it will
-only run for this amount of time).
+only run for this amount of time, and if ``forever`` is False it will
+run until disconnection).
 
 Therefore you can handle the event loop in any way you like
 instead of using ``process()``.
diff --git a/slixmpp/basexmpp.py b/slixmpp/basexmpp.py
index f59c5c00..8419a676 100644
--- a/slixmpp/basexmpp.py
+++ b/slixmpp/basexmpp.py
@@ -203,9 +203,9 @@ class BaseXMPP(XMLStream):
             log.warning('Legacy XMPP 0.9 protocol detected.')
             self.event('legacy_protocol')
 
-    def process(self, timeout=None):
+    def process(self, *, forever=True, timeout=None):
         self.init_plugins()
-        XMLStream.process(self, timeout)
+        XMLStream.process(self, forever=forever, timeout=timeout)
 
     def init_plugins(self):
         for name in self.plugin:
diff --git a/slixmpp/xmlstream/xmlstream.py b/slixmpp/xmlstream/xmlstream.py
index 820f5586..0cd6720a 100644
--- a/slixmpp/xmlstream/xmlstream.py
+++ b/slixmpp/xmlstream/xmlstream.py
@@ -207,6 +207,9 @@ class XMLStream(asyncio.BaseProtocol):
         #: ``_xmpp-client._tcp`` service.
         self.dns_service = None
 
+        #: An asyncio Future being done when the stream is disconnected.
+        self.disconnected = asyncio.Future()
+
         self.add_event_handler('disconnected', self._remove_schedules)
         self.add_event_handler('session_start', self._start_keepalive)
 
@@ -299,7 +302,7 @@ class XMLStream(asyncio.BaseProtocol):
             self.event("connection_failed", e)
             asyncio.async(self._connect_routine())
 
-    def process(self, timeout=None):
+    def process(self, *, forever=True, timeout=None):
         """Process all the available XMPP events (receiving or sending data on the
         socket(s), calling various registered callbacks, calling expired
         timers, handling signal events, etc).  If timeout is None, this
@@ -308,10 +311,15 @@ class XMLStream(asyncio.BaseProtocol):
         """
         loop = asyncio.get_event_loop()
         if timeout is None:
-            loop.run_forever()
+            if forever:
+                loop.run_forever()
+            else:
+                loop.run_until_complete(self.disconnected)
         else:
-            future = asyncio.sleep(timeout)
-            loop.run_until_complete(future)
+            tasks = [asyncio.sleep(timeout)]
+            if not forever:
+                tasks.append(self.disconnected)
+            loop.run_until_complete(asyncio.wait(tasks))
 
     def init_parser(self):
         """init the XML parser. The parser must always be reset for each new
@@ -413,6 +421,7 @@ class XMLStream(asyncio.BaseProtocol):
         if self.transport:
             self.transport.abort()
             self.event("killed")
+            self.disconnected.set_result(True)
 
     def reconnect(self, wait=2.0):
         """Calls disconnect(), and once we are disconnected (after the timeout, or