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:
@@ -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()
|
||||||
|
|
||||||
|
@@ -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()
|
||||||
|
Reference in New Issue
Block a user