Merge pull request #2081 from minrk/recv_into

implement Socket.recv_into

Author: minrk

docs/source/conf.py

@@ -272,3 +272,8 @@
 
 # If false, no module index is generated.
 # latex_use_modindex = True
+
+linkcheck_ignore = [
+    r"https://github\.com(.*)#",  # javascript based anchors
+    r"https://github\.com/zeromq/pyzmq/(issues|commits)(.*)",  # too many links
+]

examples/recv_into/discard.py

@@ -0,0 +1,95 @@
+"""
+use recv_into with an empty buffer to discard unwanted message frames
+
+avoids unnecessary allocations for message frames that won't be used
+"""
+
+import logging
+import os
+import random
+import secrets
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from threading import Thread
+
+import zmq
+
+EMPTY = bytearray()
+
+
+def subscriber(url: str) -> None:
+    log = logging.getLogger("subscriber")
+    with zmq.Context() as ctx, ctx.socket(zmq.SUB) as sub:
+        sub.linger = 0
+        sub.connect(url)
+        sub.subscribe(b"")
+        log.info("Receiving...")
+        while True:
+            frame_0 = sub.recv_string()
+            if frame_0 == "exit":
+                log.info("Exiting...")
+                break
+            elif frame_0 == "large":
+                discarded_bytes = 0
+                discarded_frames = 0
+                while sub.rcvmore:
+                    discarded_bytes += sub.recv_into(EMPTY)
+                    discarded_frames += 1
+                log.info(
+                    "Discarding large message frames: %s, bytes: %s",
+                    discarded_frames,
+                    discarded_bytes,
+                )
+            else:
+                msg: list = [frame_0]
+                if sub.rcvmore:
+                    msg.extend(sub.recv_multipart(flags=zmq.DONTWAIT))
+                log.info("Received %s", msg)
+    log.info("Done")
+
+
+def publisher(url) -> None:
+    log = logging.getLogger("publisher")
+    choices = ["large", "small"]
+    with zmq.Context() as ctx, ctx.socket(zmq.PUB) as pub:
+        pub.linger = 1000
+        pub.bind(url)
+        time.sleep(1)
+        for i in range(10):
+            kind = random.choice(choices)
+            frames = [kind.encode()]
+            if kind == "large":
+                for _ in range(random.randint(0, 5)):
+                    chunk_size = random.randint(1024, 2048)
+                    chunk = os.urandom(chunk_size)
+                    frames.append(chunk)
+            else:
+                for _ in range(random.randint(0, 3)):
+                    chunk_size = random.randint(0, 5)
+                    chunk = secrets.token_hex(chunk_size).encode()
+                    frames.append(chunk)
+            nbytes = sum(len(chunk) for chunk in frames)
+            log.info("Sending %s: %s bytes", kind, nbytes)
+            pub.send_multipart(frames)
+            time.sleep(0.1)
+        log.info("Sending exit")
+        pub.send(b"exit")
+    log.info("Done")
+
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO)
+    with TemporaryDirectory() as td:
+        sock_path = Path(td) / "example.sock"
+        url = f"ipc://{sock_path}"
+        s_thread = Thread(
+            target=subscriber, args=(url,), daemon=True, name="subscriber"
+        )
+        s_thread.start()
+        publisher(url)
+        s_thread.join(timeout=3)
+
+
+if __name__ == "__main__":
+    main()

examples/recv_into/recv_into_array.py

@@ -0,0 +1,34 @@
+"""
+Use recv_into to receive data directly into a numpy array
+"""
+
+import numpy as np
+import numpy.testing as nt
+
+import zmq
+
+url = "inproc://test"
+
+
+def main() -> None:
+    A = (np.random.random((5, 5)) * 255).astype(dtype=np.int64)
+    B = np.empty_like(A)
+    assert not (A == B).all()
+
+    with (
+        zmq.Context() as ctx,
+        ctx.socket(zmq.PUSH) as push,
+        ctx.socket(zmq.PULL) as pull,
+    ):
+        push.bind(url)
+        pull.connect(url)
+        print("sending:\n", A)
+        push.send(A)
+        bytes_received = pull.recv_into(B)
+        print(f"received {bytes_received} bytes:\n", B)
+        assert bytes_received == A.nbytes
+        nt.assert_allclose(A, B)
+
+
+if __name__ == "__main__":
+    main()

tests/test_asyncio.py

@@ -64,6 +64,61 @@ async def test_recv(create_bound_pair):
     assert recvd == b"there"
 
 
+async def test_recv_into(create_bound_pair):
+    a, b = create_bound_pair()
+    b.rcvtimeo = 1000
+    msg = [
+        b'hello',
+        b'there world',
+        b'part 3',
+        b'rest',
+    ]
+    a.send_multipart(msg)
+
+    # default nbytes: fits in array
+    buf = bytearray(10)
+    nbytes = await b.recv_into(buf)
+    assert nbytes == len(msg[0])
+    assert buf[:nbytes] == msg[0]
+
+    # default nbytes: truncates to sizeof(buf)
+    buf = bytearray(4)
+    nbytes = await b.recv_into(buf, flags=zmq.DONTWAIT)
+    # returned nbytes is the actual received length,
+    # which indicates truncation
+    assert nbytes == len(msg[1])
+    assert buf[:] == msg[1][: len(buf)]
+
+    # specify nbytes, truncates
+    buf = bytearray(10)
+    nbytes = 4
+    nbytes_recvd = await b.recv_into(buf, nbytes=nbytes)
+    assert nbytes_recvd == len(msg[2])
+
+    # recv_into empty buffer discards everything
+    buf = bytearray(10)
+    view = memoryview(buf)[:0]
+    assert view.nbytes == 0
+    nbytes = await b.recv_into(view)
+    assert nbytes == len(msg[3])
+
+
+async def test_recv_into_bad(create_bound_pair):
+    a, b = create_bound_pair()
+    b.rcvtimeo = 1000
+
+    # bad calls
+    # make sure flags work
+    with pytest.raises(zmq.Again):
+        await b.recv_into(bytearray(5), flags=zmq.DONTWAIT)
+
+    await a.send(b'msg')
+    # negative nbytes
+    buf = bytearray(10)
+    with pytest.raises(ValueError):
+        await b.recv_into(buf, nbytes=-1)
+
+
 @mark.skipif(not hasattr(zmq, "RCVTIMEO"), reason="requires RCVTIMEO")
 async def test_recv_timeout(push_pull):
     a, b = push_pull

tests/test_mypy.py

@@ -2,14 +2,14 @@
 Test our typing with mypy
 """
 
-import os
 import sys
 from pathlib import Path
 from subprocess import PIPE, STDOUT, Popen
 
 import pytest
 
 pytest.importorskip("mypy")
+pytestmark = pytest.mark.skipif(sys.version_info < (3, 10), reason="targets 3.10")
 
 repo_root = Path(__file__).parents[1]
 
@@ -25,7 +25,9 @@ def run_mypy(*mypy_args):
     Captures output and reports it on errors
     """
     p = Popen(
-        [sys.executable, "-m", "mypy"] + list(mypy_args), stdout=PIPE, stderr=STDOUT
+        [sys.executable, "-m", "mypy", "--python-version=3.10"] + list(mypy_args),
+        stdout=PIPE,
+        stderr=STDOUT,
     )
     o, _ = p.communicate()
     out = o.decode("utf8", "replace")

tests/test_socket.py

@@ -10,6 +10,7 @@
 import sys
 import time
 import warnings
+from array import array
 from unittest import mock
 
 import pytest
@@ -455,6 +456,81 @@ def test_recv_multipart(self):
         for i in range(3):
             assert self.recv_multipart(b) == [msg]
 
+    def test_recv_into(self):
+        a, b = self.create_bound_pair()
+        if not self.green:
+            b.rcvtimeo = 1000
+        msg = [
+            b'hello',
+            b'there world',
+            b'part 3',
+            b'rest',
+        ]
+        a.send_multipart(msg)
+
+        # default nbytes: fits in array
+        # make sure itemsize > 1 is handled right
+        buf = array('Q', [0])
+        nbytes = b.recv_into(buf)
+        assert nbytes == len(msg[0])
+        assert buf.tobytes()[:nbytes] == msg[0]
+
+        # default nbytes: truncates to sizeof(buf)
+        buf = bytearray(4)
+        nbytes = b.recv_into(buf)
+        # returned nbytes is the actual received length,
+        # which indicates truncation
+        assert nbytes == len(msg[1])
+        assert buf[:] == msg[1][: len(buf)]
+
+        # specify nbytes, truncates
+        buf = bytearray(10)
+        nbytes = 4
+        nbytes_recvd = b.recv_into(buf, nbytes=nbytes)
+        assert nbytes_recvd == len(msg[2])
+        assert buf[:nbytes] == msg[2][:nbytes]
+        # didn't recv excess bytes
+        assert buf[nbytes:] == bytearray(10 - nbytes)
+
+        # recv_into empty buffer discards everything
+        buf = bytearray(10)
+        view = memoryview(buf)[:0]
+        assert view.nbytes == 0
+        nbytes = b.recv_into(view)
+        assert nbytes == len(msg[3])
+        assert buf == bytearray(10)
+
+    def test_recv_into_bad(self):
+        a, b = self.create_bound_pair()
+        if not self.green:
+            b.rcvtimeo = 1000
+
+        # bad calls
+
+        # negative nbytes
+        buf = bytearray(10)
+        with pytest.raises(ValueError):
+            b.recv_into(buf, nbytes=-1)
+        # not contiguous
+        buf = memoryview(bytearray(10))[::2]
+        with pytest.raises(ValueError):
+            b.recv_into(buf)
+        # readonly
+        buf = memoryview(b"readonly")
+        with pytest.raises(ValueError):
+            b.recv_into(buf)
+        # too big
+        buf = bytearray(10)
+        with pytest.raises(ValueError):
+            b.recv_into(buf, nbytes=11)
+        # not memory-viewable
+        with pytest.raises(TypeError):
+            b.recv_into(pytest)
+
+        # make sure flags work
+        with pytest.raises(zmq.Again):
+            b.recv_into(bytearray(5), flags=zmq.DONTWAIT)
+
     def test_close_after_destroy(self):
         """s.close() after ctx.destroy() should be fine"""
         ctx = self.Context()

tests/zmq_test_utils.py

@@ -134,6 +134,7 @@ def create_bound_pair(
         s2.setsockopt(zmq.LINGER, 0)
         s2.connect(f'{interface}:{port}')
         self.sockets.extend([s1, s2])
+        s2.setsockopt(zmq.LINGER, 0)
         return s1, s2
 
     def ping_pong(self, s1, s2, msg):