How can I implement asyncio websockets in a class?

Question

I would like to connect to a websocket via asyncio and websockets, with a format as shown below. How would I be able to accomplish this?

from websockets import connect


class EchoWebsocket:

    def __init__(self):
        self.websocket = self._connect()

    def _connect(self):
        return connect("wss://echo.websocket.org")

    def send(self, message):
        self.websocket.send(message)

    def receive(self):
        return self.websocket.recv()

echo = EchoWebsocket()
echo.send("Hello!")
print(echo.receive())  # "Hello!"

Solution

How to write async programs?

1. You should define async funcs with async
2. You should call async funcs with await
3. You need event loop to start your async program

All other is almost same as with regular Python programs.

import asyncio
from websockets import connect


class EchoWebsocket:
    async def __aenter__(self):
        self._conn = connect("wss://echo.websocket.org")
        self.websocket = await self._conn.__aenter__()        
        return self

    async def __aexit__(self, *args, **kwargs):
        await self._conn.__aexit__(*args, **kwargs)

    async def send(self, message):
        await self.websocket.send(message)

    async def receive(self):
        return await self.websocket.recv()


async def main():
    async with EchoWebsocket() as echo:
        await echo.send("Hello!")
        print(await echo.receive())  # "Hello!"


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())

Output:

Hello!

As you see, code is almost same as you wrote.

Only difference is that websockets.connect designed to be async context manager (it uses __aenter__, __aexit__). It's necessary to release connection and will also help you to make async operations during class initialization (since we have no async version of __init__).

I advise you to organize your class same way. But if you really don't want to use context manager for some reason you can use new __await__ method to make async initialization and some other async function to release connection:

import sys
import asyncio
from websockets import connect


class EchoWebsocket:
    def __await__(self):
        # see: http://stackoverflow.com/a/33420721/1113207
        return self._async_init().__await__()

    async def _async_init(self):
        self._conn = connect("wss://echo.websocket.org")
        self.websocket = await self._conn.__aenter__()
        return self

    async def close(self):
        await self._conn.__aexit__(*sys.exc_info())

    async def send(self, message):
        await self.websocket.send(message)

    async def receive(self):
        return await self.websocket.recv()


async def main():
    echo = await EchoWebsocket()
    try:
        await echo.send("Hello!")
        print(await echo.receive())  # "Hello!"
    finally:
        await echo.close()


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())

Many examples of using websockets you can find in it's docs.

Comments

[Alex]

this closes after it received a reply. How could it continue receiving messages?

[Mikhail Gerasimov]

@FeedTheWeb just keep receiving replies inside context manager.

[HJA24]

@MikhailGerasimov I don't understand the above comment. How do you make sure that the script doesn't end after one reply? Thanks

[Mikhail Gerasimov]

@HJA24 if you want to send/receive data constantly, you usually start a server. For example, look at a code snippet here, specifically line server.serve_forever() . Concrete way to start a server depends on what you want to achieve. Popular high-level library to start servers that works with asyncio is aiohttp. It provides a way to work with many protocols including websockets.

[HJA24]

@MikhailGerasimov hi, I am following the docs of a different source where they don't use server.serve.forever(), so I am a little bit confused. Can you take a look at my question on SO stackoverflow.com/questions/61216022/…