Mike0021/zonos2

1from __future__ import annotations2 3from typing import Callable, Dict, Generic, TypeVar4 5import msgpack6import zmq7import zmq.asyncio8 9T = TypeVar("T")10 11 12class ZmqPushQueue(Generic[T]):13    def __init__(14        self,15        addr: str,16        create: bool,17        encoder: Callable[[T], Dict],18    ):19        self.context = zmq.Context()20        self.socket = self.context.socket(zmq.PUSH)21        self.socket.bind(addr) if create else self.socket.connect(addr)22        self.encoder = encoder23 24    def put(self, obj: T):25        event = msgpack.packb(self.encoder(obj), use_bin_type=True)26        self.socket.send(event, copy=False)27 28    def stop(self):29        self.socket.close()30        self.context.term()31 32 33class ZmqAsyncPushQueue(Generic[T]):34    def __init__(35        self,36        addr: str,37        create: bool,38        encoder: Callable[[T], Dict],39    ):40        self.context = zmq.asyncio.Context()41        self.socket = self.context.socket(zmq.PUSH)42        self.socket.bind(addr) if create else self.socket.connect(addr)43        self.encoder = encoder44 45    async def put(self, obj: T):46        event = msgpack.packb(self.encoder(obj), use_bin_type=True)47        await self.socket.send(event, copy=False)48 49    def stop(self):50        self.socket.close()51        self.context.term()52 53 54class ZmqPullQueue(Generic[T]):55    def __init__(56        self,57        addr: str,58        create: bool,59        decoder: Callable[[Dict], T],60    ):61        self.context = zmq.Context()62        self.socket = self.context.socket(zmq.PULL)63        self.socket.bind(addr) if create else self.socket.connect(addr)64        self.decoder = decoder65 66    def get(self) -> T:67        event = self.socket.recv()68        return self.decoder(msgpack.unpackb(event, raw=False))69 70    def get_raw(self) -> bytes:71        return self.socket.recv()72 73    def decode(self, raw: bytes) -> T:74        return self.decoder(msgpack.unpackb(raw, raw=False))75 76    def empty(self) -> bool:77        return self.socket.poll(timeout=0) == 078 79    def stop(self):80        self.socket.close()81        self.context.term()82 83 84class ZmqAsyncPullQueue(Generic[T]):85    def __init__(86        self,87        addr: str,88        create: bool,89        decoder: Callable[[Dict], T],90    ):91        self.context = zmq.asyncio.Context()92        self.socket = self.context.socket(zmq.PULL)93        self.socket.bind(addr) if create else self.socket.connect(addr)94        self.decoder = decoder95 96    async def get(self) -> T:97        event = await self.socket.recv()98        return self.decoder(msgpack.unpackb(event, raw=False))99 100    def stop(self):101        self.socket.close()102        self.context.term()103 104 105class ZmqPubQueue(Generic[T]):106    def __init__(107        self,108        addr: str,109        create: bool,110        encoder: Callable[[T], Dict],111    ):112        self.context = zmq.Context()113        self.socket = self.context.socket(zmq.PUB)114        self.socket.bind(addr) if create else self.socket.connect(addr)115        self.encoder = encoder116 117    def put_raw(self, raw: bytes):118        self.socket.send(raw, copy=False)119 120    def put(self, obj: T):121        event = msgpack.packb(self.encoder(obj), use_bin_type=True)122        self.socket.send(event, copy=False)123 124    def stop(self):125        self.socket.close()126        self.context.term()127 128 129class ZmqSubQueue(Generic[T]):130    def __init__(131        self,132        addr: str,133        create: bool,134        decoder: Callable[[Dict], T],135    ):136        self.context = zmq.Context()137        self.socket = self.context.socket(zmq.SUB)138        self.socket.bind(addr) if create else self.socket.connect(addr)139        self.socket.setsockopt_string(zmq.SUBSCRIBE, "")140        self.decoder = decoder141 142    def get(self) -> T:143        event = self.socket.recv()144        return self.decoder(msgpack.unpackb(event, raw=False))145 146    def empty(self) -> bool:147        return self.socket.poll(timeout=0) == 0148 149    def stop(self):150        self.socket.close()151        self.context.term()152