Python3 で websockets の async/await を隠蔽する

asyncio が入る websockets を使用した関数が上部に波及してどこまでも async/await になってしまうのか…と思っていたところ自作ライブラリでなんとかなったのでその部分の忘備録

解決策

• デコレータを使いました

環境

• Python 3.8.5
• websockets 8.1

コード

my_library.py

import websockets
import asyncio

class EntitySpec(object):
    def __init__(self, ctx):
        self.ctx = ctx
        self.uri = "ws://{}:{}".format(
            ctx.hostname, ctx.port,
        )

    def __call__(self, func):
        async def stream(func):
            async with websockets.connect(self.uri) as ws:
                while not ws.closed:
                    try:
                        response = await ws.recv()
                        func(response)
                    except websockets.exceptions.ConnectionClosedOK:
                        self.loop.stop()
        self.loop = asyncio.get_event_loop()
        self.loop.create_task(stream(func))
        return stream

    def run(self):
        self.loop.run_forever()

class Context(object):
    def __init__(
        self,
        hostname='localhost',
        port=8765,
    ):
        self.hostname = hostname
        self.port = port        
        self.websocket = EntitySpec(self)

main.py

import my_library

api = my_library.Context(
    hostname = 'localhost',
    port = 8765,  # websockets server サンプル標準ポート
)

@api.websocket
def reciever(msg):
    print(msg)

api.websocket.run()