refactor into package

msgbus/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.1"

msgbus/client.py

@@ -0,0 +1,86 @@
+import zmq
+from threading import Semaphore
+from time import sleep, time
+
+
+class PublishSetupException(Exception):
+    pass
+
+
+class MsgbusSubClient(object):
+    def __init__(self, host, port):
+        self.host = host
+        self.port = port
+        self.ctx = None  # ZMQ context
+        self.sub_socket = None  # listener sockets
+        self.subscriptions = []
+        self.pub_socket = None  # publisher socket
+        self.lock = Semaphore(1)
+        self.connect()
+
+    def close(self):
+        if self.sub_socket:
+            self.sub_socket.close()
+        if self.pub_socket:
+            self.pub_socket.close()
+        if self.ctx:
+            self.ctx.destroy()
+
+    def connect(self):
+        if not self.ctx:
+            self.ctx = zmq.Context()
+        self.sub_socket = self.ctx.socket(zmq.SUB)
+        self.sub_socket.connect("tcp://{}:{}".format(self.host, self.port))
+
+    def sub(self, channel=None):
+        if channel is None:
+            channel = ''
+        assert type(channel) is str
+        self.sub_socket.setsockopt(zmq.SUBSCRIBE, channel.encode("utf-8"))
+        self.subscriptions.append(channel)
+
+    def unsub(self, channel):
+        assert type(channel) is str
+        if channel in self.subscriptions:
+            self.subscriptions.remove(channel)
+            self.sub_socket.setsockopt(zmq.UNSUBSCRIBE, channel.encode("utf-8"))
+
+    def recv(self, decode=True, block=True):
+        recv_args = (zmq.NOBLOCK, ) if not block else ()
+        message = self.sub_socket.recv(*recv_args)
+        channel, body = message.split(b' ', 1)
+        return channel.decode("utf-8"), (body.decode('utf-8') if decode else body)
+
+    def _setup_publish_socket(self, timeout=5):
+        start = time()
+        try:
+            self.sub("__msgbus_meta")
+            while not timeout or time() < start + timeout:
+                try:
+                    channel, message = self.recv(block=False)
+                except zmq.error.Again:
+                    sleep(0.01)
+                    continue
+                if channel != "__msgbus_meta":
+                    continue
+                meta, args = message.split(" ", 1)
+                if meta != "__my_info":
+                    continue
+                server_name, subport, subproto, pubbport, pubproto = args.split(" ")
+                remote = "tcp://{}:{}".format(self.host, subport)
+                pub_socket = self.ctx.socket(zmq.PUB)
+                pub_socket.connect(remote)
+                return pub_socket
+            raise PublishSetupException("Could not establish publisher socket")
+        finally:
+            self.unsub("__msgbus_meta")
+
+    def pub(self, channel, message, encode_msg=True, settle=True, timeout=5):
+        if encode_msg:
+            message = message.encode("utf-8")
+        if not self.pub_socket:
+            with self.lock:
+                self.pub_socket = self._setup_publish_socket(timeout)
+                if settle:
+                    sleep(1)
+        self.pub_socket.send(channel.encode("utf-8") + b' ' + message)

msgbus/pub.py

@@ -0,0 +1,50 @@
+
+def send_native(host, port, channel, messages):
+    """
+    Send some messages on a specified channel using the msgbus client
+    """
+    from contextlib import closing
+    from msgbus.client import MsgbusSubClient
+    with closing(MsgbusSubClient(host, port)) as client:
+        for message in messages:
+            client.pub(channel, message)
+
+
+def send_zmq(host, port, channel, messages):
+    """
+    Send some messages on a specified channel using a raw zmq socket.
+    Note: the native client connects to the server's publisher port and autodetects the server's subscriber port. The
+    raw zmq socket must connect directly to the subscriber port
+    """
+    import zmq
+    from time import sleep
+    with zmq.Context() as ctx:
+        forward_socket = ctx.socket(zmq.PUB)
+        forward_socket.connect("tcp://{}:{}".format(host, port))
+        sleep(1)
+        for message in messages:
+            m = "{} {}".format(channel, message).encode("utf-8")
+            print(m)
+            forward_socket.send(m)
+        sleep(1)
+        forward_socket.close()
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="send a message to a msgbus server")
+    parser.add_argument("-i", "--host", default="127.0.0.1", help="host to connect to")
+    parser.add_argument("-p", "--port", default=7003, help="port to connect to")
+    parser.add_argument("-c", "--channel", required=True, help="message channel")
+    parser.add_argument("-m", "--message", required=True, nargs="+", help="message bodies")
+    parser.add_argument("--type", default="native", choices=["native", "raw"], help="client type")
+    args = parser.parse_args()
+
+    if args.type == "native":
+        send_native(args.host, args.port, args.channel, args.message)
+    elif args.type == "raw":
+        send_zmq(args.host, args.port, args.channel, args.message)
+
+
+if __name__ == '__main__':
+    main()

msgbus/sub.py

@@ -0,0 +1,56 @@
+
+def listen_native(host, port, channels):
+    """
+    Subscribe to the given server/channels using the msgbus client
+    """
+    from contextlib import closing
+    from msgbus.client import MsgbusSubClient
+    with closing(MsgbusSubClient(host, port)) as client:
+        if channels:
+            for channel in channels:
+                client.sub(channel)
+        else:
+            client.sub()
+        while True:
+            yield "{} {}".format(*client.recv())
+
+
+def listen_zmq(host, port, channels):
+    """
+    Example subscribing to the given server/channels using a raw zeromq socket
+    """
+    import zmq
+    with zmq.Context() as ctx:
+        forward_socket = ctx.socket(zmq.SUB)
+        forward_socket.connect("tcp://{}:{}".format(host, port))
+        if channels:
+            for channel in channels:
+                forward_socket.setsockopt(zmq.SUBSCRIBE, channel.encode("utf-8"))
+        else:
+            forward_socket.setsockopt(zmq.SUBSCRIBE, b'')
+        while True:
+            transport = forward_socket.recv()
+            yield transport.decode("utf-8")
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="dump all messages from msgbus")
+    parser.add_argument("-i", "--host", default="127.0.0.1", help="host to connect to")
+    parser.add_argument("-p", "--port", default=7003, help="port to connect to")
+    parser.add_argument("-c", "--channel", nargs="+", help="dump only channels")
+    parser.add_argument("--type", default="native", choices=["native", "raw"], help="client type")
+    args = parser.parse_args()
+
+    client = None
+    if args.type == "native":
+        client = listen_native(args.host, args.port, args.channel)
+    elif args.type == "raw":
+        client = listen_zmq(args.host, args.port, args.channel)
+
+    for line in client:
+        print(line)
+
+
+if __name__ == '__main__':
+    main()

setup.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+from setuptools import setup
+
+from msgbus import __version__
+
+setup(name='msgbus',
+      version=__version__,
+      description='pubsub communication tools',
+      url='http://gitlab.xmopx.net/dave/pymsgbus',
+      author='dpedu',
+      author_email='dave@davepedu.com',
+      packages=['msgbus'],
+      entry_points={
+          "console_scripts": [
+              "mbusd = msgbus.server:main",
+              "mbuspub = msgbus.pub:main",
+              "mbussub = msgbus.sub:main",
+          ]
+      })