summaryrefslogtreecommitdiff
path: root/websocket
diff options
context:
space:
mode:
Diffstat (limited to 'websocket')
-rw-r--r--websocket/_core.py13
-rw-r--r--websocket/tests/test_websocket.py13
2 files changed, 26 insertions, 0 deletions
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()