vici: Catch Python GeneratorExit to properly cancel streamed event iteration

2015-03-18 13:59:14 +01:00
parent 288b654173
commit 90c5b48c96
2 changed files with 12 additions and 1 deletions
@@ -281,10 +281,16 @@ class SessionHandler(object):
        # issue command, and read any event messages
        packet = Packet.request(command, message)
        self.transport.send(packet)
+        exited = False
        while True:
            response = Packet.parse(self.transport.receive())
            if response.response_type == Packet.EVENT:
-                yield Message.deserialize(response.payload)
+                if not exited:
+                    try:
+                        yield Message.deserialize(response.payload)
+                    except GeneratorExit:
+                        exited = True
+                        pass
            else:
                break