diff --git a/aio_pika/message.py b/aio_pika/message.py index ff1c8eff..6c10add2 100644 --- a/aio_pika/message.py +++ b/aio_pika/message.py @@ -242,7 +242,7 @@ class Message: def __init__( self, - body: bytes, + body: Union[bytes, str], *, headers: dict = None, content_type: str = None, @@ -279,7 +279,8 @@ def __init__( """ self.__lock = False - self.body = body if isinstance(body, bytes) else bytes(body) + + self.body = Message._prepare_msg_body(body) self.body_size = len(self.body) if self.body else 0 self.headers_raw = format_headers(headers) self._headers = HeaderProxy(self.headers_raw) @@ -308,6 +309,20 @@ def headers(self): def headers(self, value: dict): self.headers_raw = format_headers(value) + @staticmethod + def _prepare_msg_body(body: Union[str, bytes]) -> bytes: + """ + If body type is bytes returns body with no changes + If body type is str converts it to bytes ith encoding = 'utf8' + If body type is not either str or bytes raises TypeError + """ + if isinstance(body, bytes): + return body + elif isinstance(body, str): + return bytes(body, 'utf8') + else: + raise TypeError("body argument should be str or bytes, {} was provided instead".format(type(body))) + @staticmethod def _as_bytes(value): if isinstance(value, bytes): diff --git a/tests/test_message.py b/tests/test_message.py index 9a59a1a8..88fb6308 100644 --- a/tests/test_message.py +++ b/tests/test_message.py @@ -5,6 +5,21 @@ import shortuuid from aio_pika import DeliveryMode, Message +import pytest + + +def test_init_message_body(): + # test body is string + msg = Message(body="test") + assert msg.body == b"test" + + # test body is bytes + msg = Message(body=b"test") + assert msg.body == b"test" + + # net not str or bytes body + with pytest.raises(TypeError): + Message(body=1) def test_message_copy():