From 7932a03378f9774b0a133ce71ce65e732b5d3994 Mon Sep 17 00:00:00 2001
From: mathieui <mathieui@mathieui.net>
Date: Sun, 31 Jan 2021 17:01:41 +0100
Subject: [PATCH] xmlstream: add a simple contextmanager for temporary events

---
 slixmpp/xmlstream/xmlstream.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

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)