From 266c6b0e79cfb10299db71950f01a9a16e7c1922 Mon Sep 17 00:00:00 2001 From: Stuart Axel Owen Date: Mon, 20 Apr 2015 04:27:19 +0000 Subject: [PATCH] Add `__iter__`, `__next__`, and `next` methods to WebSocket, allowing them to be iterated over. This means you can slice, map, filter, etc, websockets without having to roll your own generator each time. --- websocket/_core.py | 13 +++++++++++++ websocket/tests/test_websocket.py | 13 +++++++++++++ 2 files changed, 26 insertions(+) diff --git a/websocket/_core.py b/websocket/_core.py index 7d73d76..4d7ddec 100644 --- a/websocket/_core.py +++ b/websocket/_core.py @@ -154,6 +154,19 @@ class WebSocket(object): else: self.lock = NoLock() + def __iter__(self): + """ + Allow iteration over websocket, implying sequential `recv` executions. + """ + while True: + yield self.recv() + + def __next__(self): + return self.recv() + + def next(self): + return self.__next__() + def fileno(self): return self.sock.fileno() diff --git a/websocket/tests/test_websocket.py b/websocket/tests/test_websocket.py index 33f7a63..81249ec 100644 --- a/websocket/tests/test_websocket.py +++ b/websocket/tests/test_websocket.py @@ -255,6 +255,19 @@ class WebSocketTest(unittest.TestCase): data = sock.recv() self.assertEqual(data, "Hello") + @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") + def testIter(self): + count = 2 + for rsvp in ws.create_connection('ws://stream.meetup.com/2/rsvps'): + count -= 1 + if count == 0: + break + + @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") + def testNext(self): + sock = ws.create_connection('ws://stream.meetup.com/2/rsvps') + self.assertEqual(str, type(next(sock))) + def testInternalRecvStrict(self): sock = ws.WebSocket() s = sock.sock = SockMock()