xmlstream: add a simple contextmanager for temporary events

1 files changed, 14 insertions, 1 deletions

diff --git a/slixmpp/xmlstream/xmlstream.py b/slixmpp/xmlstream/xmlstream.py
index 5074aa8c..b80c55d3 100644
--- a/slixmpp/xmlstream/xmlstream.py
+++ b/slixmpp/xmlstream/xmlstream.py
@@ -30,7 +30,7 @@ import weakref
 import uuid
 
 from asyncio import iscoroutinefunction, wait, Future
-
+from contextlib import contextmanager
 import xml.etree.ElementTree as ET
 
 from slixmpp.xmlstream.asyncio import asyncio
@@ -1208,3 +1208,16 @@ class XMLStream(asyncio.BaseProtocol):
             disposable=True,
         )
         return await asyncio.wait_for(fut, timeout, loop=self.loop)
+
+    @contextmanager
+    def event_handler(self, event: str, handler: Callable):
+        """
+        Context manager that adds then removes an event handler.
+        """
+        self.add_event_handler(event, handler)
+        try:
+            yield
+        except Exception as exc:
+            raise
+        finally:
+            self.del_event_handler(event, handler)