Support for rooms

docs/index.rst

@@ -32,7 +32,7 @@ The following code example shows how to add Flask-SocketIO to a Flask applicatio
 
 The ``init_app()`` style of initialization is also supported. Note the way the web server is started. The ``socketio.run()`` function encapsulates the start up of the gevent web server.
 
-Note that the Werkzeug development web server cannot be used with this extension, but its debugger and reloader modules are activated when the application is in debug mode.
+The Werkzeug development web server cannot be used with this extension, but its debugger and reloader modules are activated when the application is in debug mode.
 
 Receiving Messages
 ------------------
@@ -109,15 +109,45 @@ SocketIO supports acknowledgement callbacks that confirm that a message was rece
     def handle_my_custom_event(json):
         emit('my response', json, callback=ack)
 
-When using callbacks the Javascript client receives a callback function to invoke upon receipt of the message. When client calls the callback function the server is notified and invokes the corresponding server-side callback. The client-side can pass arguments in the callback function, which are transferred to the server and given to the server-side callback..
+When using callbacks the Javascript client receives a callback function to invoke upon receipt of the message. When client calls the callback function the server is notified and invokes the corresponding server-side callback. The client-side can pass arguments in the callback function, which are transferred to the server and given to the server-side callback.
 
 Another very useful feature of SocketIO is the broadcasting of messages. Flask-SocketIO supports this feature with the ``broadcast=True`` optional argument to ``send()`` and ``emit()``::
 
     @socketio.on('my event')
     def handle_my_custom_event(data):
         emit('my response', data, broadcast=True)
 
-When a message is sent with the broadcast option enabled all clients connected to the namespace receive it, including the sender. When namespaces are not used the clients connected to the global namespace receive the message.
+When a message is sent with the broadcast option enabled all clients connected to the namespace receive it, including the sender. When namespaces are not used the clients connected to the global namespace receive the message. Callbacks are not invoked for broadcast messages.
+
+Sometimes the server needs to be the originator of a message. This can be useful to send a notification to clients of an event that originated in the server. The ``socketio.send()`` and ``socketio.emit()`` methods can be used to broadcast to all connected clients::
+
+    def some_function():
+        socketio.emit('some event', {'data': 42})
+
+Note that in this usage the ``broadcast=True`` argument does not need to be specified.
+
+Rooms
+-----
+
+For many applications it is necessary to group subsets of users to send messages. The best example is a chat application with multiple rooms. Users should receive messages from the room or rooms they are in, but not from other rooms where other users are. Flask-SocketIO supports this concept of rooms through the ``join_room()`` and ``leave_room()`` functions::
+
+    from flask.ext.socketio import join_room, leave_room
+
+    @socketio.on('join')
+    def on_join(data):
+        username = data['username']
+        room = data['room']
+        join_room(room)
+        send(username + ' has entered the room.', room=room)
+
+    @socketio.on('leave')
+    def on_leave(data):
+        username = data['username']
+        room = data['room']
+        leave_room(room)
+        send(username + ' has left the room.', room=room)
+
+The ``send()`` and ``emit()`` functions accept an optional ``room`` argument that cause the message to be sent to all the clients that are in the given room. A given client can join multiple rooms if desired. When a client disconnects it is removed from any room it was in.
 
 Connection Events
 -----------------

example/app.py

@@ -1,5 +1,5 @@
-from flask import Flask, render_template, session
-from flask.ext.socketio import SocketIO, emit
+from flask import Flask, render_template, session, request
+from flask.ext.socketio import SocketIO, emit, join_room, leave_room
 
 app = Flask(__name__)
 app.debug=True
@@ -27,6 +27,32 @@ def test_message(message):
          broadcast=True)
 
 
+@socketio.on('join', namespace='/test')
+def join(message):
+    join_room(message['room'])
+    session['receive_count'] = session.get('receive_count', 0) + 1
+    emit('my response',
+         {'data': 'In rooms: ' + ', '.join(request.namespace.rooms),
+          'count': session['receive_count']})
+
+
+@socketio.on('leave', namespace='/test')
+def leave(message):
+    leave_room(message['room'])
+    session['receive_count'] = session.get('receive_count', 0) + 1
+    emit('my response',
+         {'data': 'In rooms: ' + ', '.join(request.namespace.rooms),
+          'count': session['receive_count']})
+
+
+@socketio.on('my room event', namespace='/test')
+def send_room_message(message):
+    session['receive_count'] = session.get('receive_count', 0) + 1
+    emit('my response',
+        {'data': message['data'], 'count': session['receive_count']},
+        room=message['room'])
+
+
 @socketio.on('connect', namespace='/test')
 def test_connect():
     emit('my response', {'data': 'Connected', 'count': 0})
@@ -39,3 +65,4 @@ def test_disconnect():
 
 if __name__ == '__main__':
     socketio.run(app)
+

example/templates/index.html

@@ -8,7 +8,7 @@
         $(document).ready(function(){
             var socket = io.connect('http://' + document.domain + ':' + location.port + '/test');
             socket.on('my response', function(msg) {
-                $('#log').append('<p>Received #' + msg.count + ': ' + msg.data + '</p>');
+                $('#log').append('<br>Received #' + msg.count + ': ' + msg.data);
             });
             $('form#emit').submit(function(event) {
                 socket.emit('my event', {data: $('#emit_data').val()});
@@ -18,28 +18,45 @@
                 socket.emit('my broadcast event', {data: $('#broadcast_data').val()});
                 return false;
             });
+            $('form#join').submit(function(event) {
+                socket.emit('join', {room: $('#join_room').val()});
+                return false;
+            });
+            $('form#leave').submit(function(event) {
+                socket.emit('leave', {room: $('#leave_room').val()});
+                return false;
+            });
+            $('form#send_room').submit(function(event) {
+                socket.emit('my room event', {room: $('#room_name').val(), data: $('#room_data').val()});
+                return false;
+            });
         });
     </script>
 </head>
 <body>
     <h1>Flask-SocketIO Test</h1>
     <h2>Send:</h2>
-    <table>
-        <tr>
-            <td>
-                <form id="emit" method='POST' action='#'>
-                    <textarea name="emit_data" id="emit_data"></textarea>
-                    <div><input type="submit" value="Emit"></div>
-                </form>
-            </td>
-            <td>
-                <form id="broadcast" method='POST' action='#'>
-                    <textarea name="broadcast_data" id="broadcast_data"></textarea>
-                    <div><input type="submit" value="Broadcast"></div>
-                </form>
-            </td>
-        </tr>
-    </table>
+    <form id="emit" method='POST' action='#'>
+        <input type="text" name="emit_data" id="emit_data" placeholder="Message">
+        <input type="submit" value="Echo"></div>
+    </form>
+    <form id="broadcast" method='POST' action='#'>
+        <input type="text" name="broadcast_data" id="broadcast_data" placeholder="Message">
+        <input type="submit" value="Broadcast">
+    </form>
+    <form id="join" method='POST' action='#'>
+        <input type="text" name="join_room" id="join_room" placeholder="Room Name">
+        <input type="submit" value="Join Room">
+    </form>
+    <form id="leave" method='POST' action='#'>
+        <input type="text" name="leave_room" id="leave_room" placeholder="Room Name">
+        <input type="submit" value="Leave Room">
+    </form>
+    <form id="send_room" method='POST' action='#'>
+        <input type="text" name="room_name" id="room_name" placeholder="Room Name">
+        <input type="text" name="room_data" id="room_data" placeholder="Message">
+        <input type="submit" value="Send to Room">
+    </form>
     <h2>Receive:</h2>
     <div id="log"></div>
 </body>

flask_socketio/__init__.py

@@ -30,6 +30,7 @@ def __init__(self, app=None):
         if app:
             self.init_app(app)
         self.messages = {}
+        self.rooms = {}
 
     def init_app(self, app):
         app.wsgi_app = SocketIOMiddleware(app, self)
@@ -40,60 +41,62 @@ class GenericNamespace(base_namespace):
             base_emit = base_namespace.emit
             base_send = base_namespace.send
 
+            def initialize(self):
+                self.rooms = set()
+
             def process_event(self, packet):
                 message = packet['name']
                 args = packet['args']
                 app = self.request
-                return self.socketio.dispatch_message(app, self, message, args)
+                return self.socketio._dispatch_message(app, self, message, args)
+
+            def join_room(self, room):
+                if self.socketio._join_room(self, room):
+                    self.rooms.add(room)
+
+            def leave_room(self, room):
+                if self.socketio._leave_room(self, room):
+                    self.rooms.remove(room)
 
             def recv_connect(self):
                 ret = super(GenericNamespace, self).recv_connect()
                 app = self.request
-                self.socketio.dispatch_message(app, self, 'connect')
+                self.socketio._dispatch_message(app, self, 'connect')
                 return ret
 
             def recv_disconnect(self):
                 app = self.request
-                self.socketio.dispatch_message(app, self, 'disconnect')
+                self.socketio._dispatch_message(app, self, 'disconnect')
+                for room in self.rooms.copy():
+                    self.leave_room(room)
                 return super(GenericNamespace, self).recv_disconnect()
 
             def recv_message(self, data):
                 app = self.request
-                return self.socketio.dispatch_message(app, self, 'message', [data])
+                return self.socketio._dispatch_message(app, self, 'message', [data])
 
             def recv_json(self, data):
                 app = self.request
-                return self.socketio.dispatch_message(app, self, 'json', [data])
+                return self.socketio._dispatch_message(app, self, 'json', [data])
 
             def emit(self, event, *args, **kwargs):
                 ns_name = kwargs.pop('namespace', None)
                 broadcast = kwargs.pop('broadcast', False)
-                if broadcast:
+                room = kwargs.pop('room', None)
+                if broadcast or room:
                     if ns_name is None:
                         ns_name = self.ns_name
-                    callback = kwargs.pop('callback', None)
-                    ret = None
-                    for sessid, socket in self.socket.server.sockets.items():
-                        if socket == self.socket:
-                            ret = self.base_emit(event, *args, callback=callback, **kwargs)
-                        else:
-                            socket[ns_name].base_emit(event, *args, **kwargs)
-                    return ret
+                    return self.socketio.emit(event, *args, namespace=ns_name, room=room)
                 if ns_name is None:
                     return self.base_emit(event, *args, **kwargs)
                 return request.namespace.socket[ns_name].base_emit(event, *args, **kwargs)
 
-            def send(self, message, json=False, ns_name=None, callback=None, broadcast=False):
-                if broadcast:
+            def send(self, message, json=False, ns_name=None, callback=None,
+                     broadcast=False, room=None):
+                if broadcast or room:
                     if ns_name is None:
                         ns_name = self.ns_name
-                    ret = None
-                    for sessid, socket in self.socket.server.sockets.items():
-                        if socket == request.namespace.socket:
-                            ret = self.base_send(message, json, callback=callback)
-                        else:
-                            socket[ns_name].base_send(message, json)
-                    return ret
+                    return self.socketio.send(message, json, ns_name, room)
                 if ns_name is None:
                     return request.namespace.base_send(message, json, callback)
                 return request.namespace.socket[ns_name].base_send(message, json, callback)
@@ -103,7 +106,7 @@ def send(self, message, json=False, ns_name=None, callback=None, broadcast=False
             namespaces[ns_name] = GenericNamespace
         return namespaces
 
-    def dispatch_message(self, app, namespace, message, args=[]):
+    def _dispatch_message(self, app, namespace, message, args=[]):
         if namespace.ns_name not in self.messages:
             return
         if message not in self.messages[namespace.ns_name]:
@@ -117,6 +120,29 @@ def dispatch_message(self, app, namespace, message, args=[]):
                 namespace.session[k] = v
             return ret
 
+    def _join_room(self, namespace, room):
+        if namespace.ns_name not in self.rooms:
+            self.rooms[namespace.ns_name] = {}
+        if room not in self.rooms[namespace.ns_name]:
+            self.rooms[namespace.ns_name][room] = set()
+        if namespace not in self.rooms[namespace.ns_name][room]:
+            self.rooms[namespace.ns_name][room].add(namespace)
+            return True
+        return False
+
+    def _leave_room(self, namespace, room):
+        if namespace.ns_name in self.rooms:
+            if room in self.rooms[namespace.ns_name]:
+                if namespace in self.rooms[namespace.ns_name][room]:
+                    self.rooms[namespace.ns_name][room].remove(namespace)
+                    if len(self.rooms[namespace.ns_name][room]) == 0:
+                        del self.rooms[namespace.ns_name][room]
+                        if len(self.rooms[namespace.ns_name]) == 0:
+                            del self.rooms[namespace.ns_name]
+
+                    return True
+        return False
+
     def on_message(self, message, handler, **options):
         ns_name = options.pop('namespace', '')
         if ns_name not in self.messages:
@@ -129,6 +155,31 @@ def decorator(f):
             return f
         return decorator
 
+    def emit(self, event, *args, **kwargs):
+        ns_name = kwargs.pop('namespace', None)
+        if ns_name is None:
+            ns_name = ''
+        room = kwargs.pop('room', None)
+        if room:
+            for client in self.rooms.get(ns_name, {}).get(room, set()):
+                client.base_emit(event, *args, **kwargs)
+        else:
+            for sessid, socket in self.server.sockets.items():
+                if socket.active_ns.get(ns_name):
+                    socket[ns_name].base_emit(event, *args, **kwargs)
+
+    def send(self, message, json=False, namespace=None, room=None):
+        ns_name = namespace
+        if ns_name is None:
+            ns_name = ''
+        if room:
+            for client in self.rooms.get(ns_name, {}).get(room, set()):
+                client.base_send(message, json)
+        else:
+            for sessid, socket in self.server.sockets.items():
+                if socket.active_ns.get(ns_name):
+                    socket[ns_name].base_send(message, json)
+
     def run(self, app, host=None, port=None):
         if host is None:
             host = '127.0.0.1'
@@ -138,14 +189,14 @@ def run(self, app, host=None, port=None):
                 port = int(server_name.rsplit(':', 1)[1])
             else:
                 port = 5000
+        self.server = SocketIOServer((host, port), app.wsgi_app, resource='socket.io')
         if app.debug:
             @run_with_reloader
             def run_server():
-                server = SocketIOServer((host, port), app.wsgi_app, resource='socket.io')
-                server.serve_forever()
+                self.server.serve_forever()
             run_server()
         else:
-            SocketIOServer((host, port), app.wsgi_app, resource='socket.io').serve_forever()
+            self.server.serve_forever()
 
     def test_client(self, app, namespace=None):
         return SocketIOTestClient(app, self, namespace)
@@ -154,9 +205,16 @@ def test_client(self, app, namespace=None):
 def emit(event, *args, **kwargs):
     return request.namespace.emit(event, *args, **kwargs)
 
+def send(message, json=False, namespace=None, callback=None, broadcast=False, room=None):
+    return request.namespace.send(message, json, namespace, callback, broadcast, room)
+
+
+def join_room(room):
+    return request.namespace.join_room(room)
+
 
-def send(message, json=False, namespace=None, callback=None, broadcast=False):
-    return request.namespace.send(message, json, namespace, callback, broadcast)
+def leave_room(room):
+    return request.namespace.leave_room(room)
 
 
 def error(error_name, error_message, msg_id=None, quiet=False):