Merge branch 'master' of https://github.com/StuartAxelOwen/websocket-client into StuartAxelOwen-master

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()