ログインが必要なWebsocketを扱えるライブラリつくった

先日書いた記事で必要なことを、さくっとできるようにライブラリにしてPyPiに登録してみました。

github.com

どんなライブラリ

requestsを使って取得したcookieをシームレスにwebsocket-clientに引き渡すラッパーライブラリです。connectを実行する際に、その時点で保持しているrequestsのcookieをすべて文字列に変換して引き渡すことで、セッション情報を共有しながらwebsocketの接続をできるようにしています。

使用例

インストールはpipだけでできます。

$ pip install wsrequests

簡単な使い方はこんな感じになります。

from wsrequests import WsRequests

wsr = WsRequests()

# login django
wsr.get('http://localhost:3000/login')
wsr.post(
    'http://localhost:3000/login',
    data={
        'username': 'your name',
        'password': 'your password',
        'csrfmiddlewaretoken': wsr.cookies['csrftoken'],
        'next': '/',
    }
)

WsRequestsをインスタンス化した時点で内部的にはrequests.session()でセッションが生成されます。getやpostはそのままそのセッションインスタンスのメソッドとして呼び出されるので、インターフェースはrequestsそのままです。

ここまで問題なければセッションはログイン状態になっていますので、以下のようにWebSocketの処理ができます。

# connect websocket
wsr.connect('ws://localhost:3000/ws/room')
wsr.send_message({'message': 'ping'})
wsr.receive_message()
wsr.disconnect()

内部的な話

WsRequestsの説明にも書いてありますが、こいつはrequestsの薄いラッパーです。getやpostを始めとした大半のメソッドはrequestsを使うようにし、WebSocket関連の処理だけwebsocket-clientを使ったメソッドを呼び出せるようにすれば要件を満たせます。

ということで、実装はこんな感じになっています。

class WsRequests:

    def __init__(self, proxy_url=None, proxy_port=None, proxy_username=None, proxy_password=None, logger=None):
        :
        :
        self.http_client = self._get_http_client()
        :

    def __getattr__(self, attr_name):
        """
        Access for not defined attribute dispatch to requests.
        :param str attr_name:
        :return:
        """
        return getattr(self.http_client, attr_name)
        :
    def _get_http_client(self):
        """
        get requests session instance
        :return:
        """
        client = requests.session()
        if self.is_valid_proxy:
            client.trust_env = False
            client.proxies.update({
                'http': self._get_proxy_url(),
                'https': self._get_proxy_url(),
            })
            for proxy in client.proxies:
                self.logger.debug('use proxy [{}]'.format(proxy))
        return client

インスタンス化した時点でself._get_http_clientを呼び出し、その中でrequests.session()で取得したインスタンスを戻し、self.http_clientにセットしています。そして、ポイントは__getattr__をオーバライドし、属性の呼び出しはself.http_clientにディスパッチしています。こうすることでWsRequests自身が実装しているメソッド以外はrequests.session()のものが呼び出されるので、ほとんどrequestsのインターフェースを保った形になっています。

最後に

とりあえず、必要な機能は実装できていますが今後要件が増えたら随時機能追加をしていこうと思います。

github.com