Supports for multibytes and multidimensional memoryview (#4890)

CHANGES/4890.bugfix

@@ -0,0 +1 @@
+Fix incorrect size calculation for memoryviews

CONTRIBUTORS.txt

@@ -130,6 +130,7 @@ Igor Davydenko
 Igor Mozharovsky
 Igor Pavlov
 Ilya Chichak
+Ilya Gruzinov
 Ingmar Steen
 Ivan Larin
 Jacob Champion

aiohttp/http_writer.py

@@ -78,6 +78,11 @@ async def write(self, chunk: bytes,
  if self._on_chunk_sent is not None:
  await self._on_chunk_sent(chunk)
 
+ if isinstance(chunk, memoryview):
+ if chunk.nbytes != len(chunk):
+ # just reshape it
+ chunk = chunk.cast('c')
+
  if self._compress is not None:
  chunk = self._compress.compress(chunk)
  if not chunk:

aiohttp/payload.py

@@ -215,7 +215,10 @@ def __init__(self,
 
  super().__init__(value, *args, **kwargs)
 
- self._size = len(value)
+ if isinstance(value, memoryview):
+ self._size = value.nbytes
+ else:
+ self._size = len(value)
 
  if self._size > TOO_LARGE_BYTES_BODY:
  if PY_36:

tests/test_http_writer.py

@@ -1,4 +1,5 @@
 # Tests for aiohttp/http_writer.py
+import array
 from unittest import mock
 
 import pytest
@@ -152,6 +153,88 @@ async def test_write_payload_deflate_and_chunked(
  )
  assert thing == buf
 
+async def test_write_payload_bytes_memoryview(
+ buf,
+ protocol,
+ transport,
+ loop):
+
+ msg = http.StreamWriter(protocol, loop)
+
+ mv = memoryview(b"abcd")
+
+ await msg.write(mv)
+ await msg.write_eof()
+
+ thing = b"abcd"
+ assert thing == buf
+
+
+async def test_write_payload_short_ints_memoryview(
+ buf,
+ protocol,
+ transport,
+ loop):
+ msg = http.StreamWriter(protocol, loop)
+ msg.enable_chunking()
+
+ payload = memoryview(array.array("H", [65, 66, 67]))
+
+ await msg.write(payload)
+ await msg.write_eof()
+
+ endians = (
+ (
+ b"6\r\n"
+ b"\x00A\x00B\x00C\r\n"
+ b'0\r\n\r\n'
+ ),
+ (
+ b"6\r\n"
+ b"A\x00B\x00C\x00\r\n"
+ b"0\r\n\r\n"
+ )
+ )
+ assert buf in endians
+
+
+async def test_write_payload_2d_shape_memoryview(
+ buf,
+ protocol,
+ transport,
+ loop):
+ msg = http.StreamWriter(protocol, loop)
+ msg.enable_chunking()
+
+ mv = memoryview(b"ABCDEF")
+ payload = mv.cast("c", [3, 2])
+
+ await msg.write(payload)
+ await msg.write_eof()
+
+ thing = (
+ b"6\r\n"
+ b"ABCDEF\r\n"
+ b"0\r\n\r\n"
+ )
+ assert thing == buf
+
+async def test_write_payload_slicing_long_memoryview(
+ buf,
+ protocol,
+ transport,
+ loop):
+ msg = http.StreamWriter(protocol, loop)
+ msg.length = 4
+
+ mv = memoryview(b"ABCDEF")
+ payload = mv.cast("c", [3, 2])
+
+ await msg.write(payload)
+ await msg.write_eof()
+
+ thing = b"ABCD"
+ assert thing == buf
 
 async def test_write_drain(protocol, transport, loop) -> None:
  msg = http.StreamWriter(protocol, loop)

tests/test_payload.py

@@ -1,3 +1,4 @@
+import array
 from io import StringIO
 
 import pytest
@@ -65,6 +66,12 @@ def test_bytes_payload_bad_type() -> None:
  payload.BytesPayload(object())
 
 
+def test_bytes_payload_memoryview_correct_size() -> None:
+ mv = memoryview(array.array("H", [1, 2, 3]))
+ p = payload.BytesPayload(mv)
+ assert p.size == 6
+
+
 def test_string_payload() -> None:
  p = payload.StringPayload('test')
  assert p.encoding == 'utf-8'