pico.chat/pico.py

import asyncio
import collections
import datetime
import http
import json
import functools
from pathlib import Path

import websockets


rooms = collections.defaultdict(dict)


async def send_json_many(targets, **data):
    for websocket in list(targets):
        await send_json(websocket, **data)


async def send_json(websocket, **data):
    try:
        await websocket.send(json.dumps(data))
    except websockets.exceptions.ConnectionClosed:
        pass


async def recv_json(websocket):
    try:
        return json.loads(await websocket.recv())
    except websockets.exceptions.ConnectionClosed:
        return {'action': 'logout'}
    except json.decoder.JSONDecodeError:
        return {}


async def core(websocket, path, server_name):
    sockets = rooms[path]
    while True:
        data = await recv_json(websocket)
        ts = datetime.datetime.now().isoformat()
        reply = functools.partial(send_json, websocket=websocket, ts=ts)
        error = functools.partial(reply, kind='error')
        broadcast = functools.partial(send_json_many, targets=sockets, ts=ts)

        if 'action' not in data:
            await error(info='Message without action is invalid')

        elif websocket not in sockets and data['action'] not in {'login', 'logout'}:
            await error(info='Not logged in')

        elif data['action'] == 'login':
            username = data['username']
            if not username:
                await error(info='Invalid username')
            elif username in sockets.values():
                await error(info='Username taken')
            else:
                sockets[websocket] = username
                await reply(kind='update', username=username)
                await broadcast(kind='update', users=list(sockets.values()), info=f'{username} joined')

        elif data['action'] == 'post':
            text = data['text']
            if not text:
                continue
            await broadcast(kind='post', source=sockets[websocket], text=text)

        elif data['action'] == 'logout':
            if websocket in sockets:
                username = sockets.pop(websocket)
            await broadcast(kind='update', users=list(sockets.values()), info=f'{username} left')
            break


async def serve_html(path, request_headers):
    if request_headers.get('Upgrade') != 'websocket':

        document = Path(__file__, '..', Path(path).name).resolve()
        if not document.is_file():
            document = Path(__file__, '..', 'pico.html').resolve()
            content_type = 'text/html; charset=utf-8'
        elif path.endswith('.js'):
            content_type = 'application/javascript; charset=utf-8'
        elif path.endswith('.css'):
            content_type = 'text/css; charset=utf-8'
        else:
            content_type = 'text/plain; charset=utf-8'

        return (
            http.HTTPStatus.OK,
            [('Content-Type', content_type)],
            document.read_bytes(),
        )


async def start_server(host, port, server_name):
    bound_core = functools.partial(core, server_name=server_name)
    return await websockets.serve(bound_core, host, port, process_request=serve_html)


if __name__ == '__main__':
    host, port = 'localhost', 9753
    loop = asyncio.get_event_loop()
    loop.run_until_complete(start_server(host, port, 'PicoChat'))
    print(f'Running on {host}:{port}')
    loop.run_forever()