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.

This commit is contained in:
Stuart Axel Owen
2015-04-20 04:27:19 +00:00
parent b96a2e8fa7
commit 266c6b0e79
2 changed files with 26 additions and 0 deletions

View File

@@ -154,6 +154,19 @@ class WebSocket(object):
else: else:
self.lock = NoLock() 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): def fileno(self):
return self.sock.fileno() return self.sock.fileno()

View File

@@ -255,6 +255,19 @@ class WebSocketTest(unittest.TestCase):
data = sock.recv() data = sock.recv()
self.assertEqual(data, "Hello") 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): def testInternalRecvStrict(self):
sock = ws.WebSocket() sock = ws.WebSocket()
s = sock.sock = SockMock() s = sock.sock = SockMock()