Skip to content

Commit

Permalink
Add json convenience methods to websocket message objects
Browse files Browse the repository at this point in the history
  • Loading branch information
sseg committed Jun 2, 2016
1 parent c059e0b commit d55740d
Show file tree
Hide file tree
Showing 5 changed files with 543 additions and 422 deletions.
9 changes: 9 additions & 0 deletions aiohttp/web_ws.py
Original file line number Diff line number Diff line change
Expand Up @@ -285,6 +285,15 @@ def receive_bytes(self):
msg.data))
return msg.data

@asyncio.coroutine
def receive_json(self):
msg = yield from self.receive()
if msg.tp != MsgType.text:
raise TypeError(
"Received message {}:{!r} is not str".format(msg.tp, msg.data)
)
return msg.json()

def write(self, data):
raise RuntimeError("Cannot call .write() for websocket")

Expand Down
11 changes: 10 additions & 1 deletion aiohttp/websocket.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import binascii
import collections
import hashlib
import json
import os
import random
import sys
Expand Down Expand Up @@ -57,7 +58,6 @@
hdrs.SEC_WEBSOCKET_KEY,
hdrs.SEC_WEBSOCKET_PROTOCOL)

Message = collections.namedtuple('Message', ['tp', 'data', 'extra'])

UNPACK_LEN2 = Struct('!H').unpack_from
UNPACK_LEN3 = Struct('!Q').unpack_from
Expand All @@ -69,6 +69,15 @@
MSG_SIZE = 2 ** 14


MessageBase = collections.namedtuple('Message', ['tp', 'data', 'extra'])


class Message(MessageBase):
def json(self, *, loads=json.loads):
"""Return parsed JSON data."""
return loads(self.data)


class WebSocketError(Exception):
"""WebSocket protocol parser error."""

Expand Down
24 changes: 20 additions & 4 deletions docs/web_reference.rst
Original file line number Diff line number Diff line change
Expand Up @@ -260,9 +260,9 @@ like one using :meth:`Request.copy`.

async def json(self, *, loads=json.loads):
body = await self.text()
return loader(body)
return loads(body)

:param callable loader: any :term:`callable` that accepts
:param callable loads: any :term:`callable` that accepts
:class:`str` and returns :class:`dict`
with parsed JSON (:func:`json.loads` by
default).
Expand Down Expand Up @@ -894,7 +894,7 @@ WebSocketResponse

.. coroutinemethod:: receive_str()

A :ref:`coroutine<coroutine>` that calls :meth:`receive_mgs` but
A :ref:`coroutine<coroutine>` that calls :meth:`receive` but
also asserts the message type is
:const:`~aiohttp.websocket.MSG_TEXT`.

Expand All @@ -904,14 +904,30 @@ WebSocketResponse

.. coroutinemethod:: receive_bytes()

A :ref:`coroutine<coroutine>` that calls :meth:`receive_mgs` but
A :ref:`coroutine<coroutine>` that calls :meth:`receive` but
also asserts the message type is
:const:`~aiohttp.websocket.MSG_BINARY`.

:return bytes: peer's message content.

:raise TypeError: if message is :const:`~aiohttp.websocket.MSG_TEXT`.

.. coroutinemethod:: recieve_json(*, loads=json.loads)

A :ref:`coroutine<coroutine>` that calls :meth:`receive`, asserts the
message type is :const:`~aiohttp.websocket.MSG_TEXT`, and loads the JSON
string to a Python dict.

:param callable loads: any :term:`callable` that accepts
:class:`str` and returns :class:`dict`
with parsed JSON (:func:`json.loads` by
default).

:return dict: loaded JSON content

:raise TypeError: if message is :const:`~aiohttp.websocket.MSG_BINARY`.
:raise ValueError: if message is not valid JSON.


.. versionadded:: 0.14

Expand Down
Loading

0 comments on commit d55740d

Please sign in to comment.