winwin
winwin

Reputation: 1834

How to properly use asyncio in this context?

I have an idea of using an event-loop in a websocket-based server (based on websockets library), but I'm struggling to do so, since I really don't get the idea of how to use asyncio and write your own coroutines.

# let us define a dictionary of callbacks awaiting to be triggered
CALLBACKS = dict()

# let us have an infinite loop, which reads new messages from a websocket connection
async def websocket_connection(ws, path):
     async for msg in ws:
         # we received a message
         msg = json.loads(msg)
         msg_id = msg['id']
         msg_data = msg['data']
         # we check, if there is a callback waiting to be triggered with such id
         if msg_id in CALLBACKS:
             # if such id exists, we get the callback function
             callback = CALLBACKS[msg_id]
             # we delete the callback from the dictionary
             del CALLBACKS[msg_id]
             # we call it passing the data from the message in there
             callback(msg_data)
         else:
             # we don't have such id, throw some error
             raise Exception("bad id " + msg_id)


# this is the function which submits a message
async def submit(ws, data):
    # it generates a random request id
    from uuid import uuid4
    request_id = str(uuid4())
    # registers a callback, which will simply return the received data
    CALLBACKS[request_id] = lambda data:  # <- this is the part I don't know what it should really do to be able to return the value from the await
    # after it encodes the message as json
    msg = json.dumps({'id': request_id, 'data': data})
    # and sends via the websocket
    await ws.send(msg)

# and let us have the actual function, which will be the actual script
async def websocket_script(ws):
    # this is what I would like to be able to do in the end.
    # pass in commands to the websocket and await their completion
    sum = int(await submit(ws, {"eval": "2+2"}))
    division = float(await submit(ws, {"eval": "3/2"}))

websocket_connection and websocket_script would need to be going side by side for it to work. I suppose gather or some other async function would work. I would really like to get rid of callbacks, since this is the initial purpose of using asyncio in the first place.

How could this be done? It seems like a job for asyncio.

Upvotes: 1

Views: 180

Answers (1)

Mikhail Gerasimov
Mikhail Gerasimov

Reputation: 39586

I think what you search is asyncio.Future if I understand you correctly.

FUTURES = dict()

async def websocket_connection(ws, path):
    # ...
    fut = FUTURES[msg_id]
    fut.set_result(msg_data)  # will "unblock" await fut below
    del FUTURES[msg_id]
    # ...


async def submit(ws, data):
    from uuid import uuid4
    request_id = str(uuid4())

    fut = asyncio.Future()
    FUTURES[request_id] = fut

    msg = json.dumps({'id': request_id, 'data': data})    
    await ws.send(msg)
    # ...
    
    return (await fut)


async def websocket_script(ws):
    sum = int(await submit(ws, {"eval": "2+2"}))
    division = float(await submit(ws, {"eval": "3/2"}))

Upvotes: 1

Related Questions