ZMQ classed updates to support IPC protocol

Author: mrakitin

src/bluesky/callbacks/zmq.py

@@ -15,7 +15,9 @@
 import asyncio
 import copy
 import pickle
+import uuid
 import warnings
+from enum import Enum
 
 from ..run_engine import Dispatcher, DocumentNames
 
@@ -26,6 +28,11 @@ class Bluesky0MQDecodeError(Exception):
     ...
 
 
+class Protocols(Enum):
+    TCP = "tcp"
+    IPC = "ipc"
+
+
 class Publisher:
     """
     A callback that publishes documents to a 0MQ proxy.
@@ -58,7 +65,7 @@ class Publisher:
     >>> RE.subscribe(publisher)
     """
 
-    def __init__(self, address, *, prefix=b"", RE=None, zmq=None, serializer=pickle.dumps):
+    def __init__(self, address, *, prefix=b"", RE=None, zmq=None, serializer=pickle.dumps, protocol=Protocols.TCP):
         if RE is not None:
             warnings.warn(  # noqa: B028
                 "The RE argument to Publisher is deprecated and "
@@ -73,11 +80,15 @@ def __init__(self, address, *, prefix=b"", RE=None, zmq=None, serializer=pickle.
             raise ValueError(f"prefix {prefix!r} may not contain b' '")
         if zmq is None:
             import zmq
-        if isinstance(address, str):
-            address = address.split(":", maxsplit=1)
-        self.address = (address[0], int(address[1]))
+
+        if isinstance(address, tuple):
+            url = f"{protocol.value}://{address[0]}:{address[1]}"
+        else:
+            url = f"{protocol.value}://address"
+
+        self.address = url
         self.RE = RE
-        url = "tcp://%s:%d" % self.address
+
         self._prefix = bytes(prefix)
         self._context = zmq.Context()
         self._socket = self._context.socket(zmq.PUB)
@@ -116,9 +127,9 @@ class Proxy:
 
     Attributes
     ----------
-    in_port : int
+    in_port : int or str
         Port that RunEngines should broadcast to.
-    out_port : int
+    out_port : int or str
         Port that subscribers should subscribe to.
     closed : boolean
         True if the Proxy has already been started and subsequently
@@ -146,28 +157,41 @@ class Proxy:
     >>> proxy.start()  # runs until interrupted
     """
 
-    def __init__(self, in_port=None, out_port=None, *, zmq=None):
+    def __init__(self, in_port=None, out_port=None, *, zmq=None, protocol=Protocols.TCP):
         if zmq is None:
             import zmq
         self.zmq = zmq
         self.closed = False
+        self.input_path = None
+        self.output_path = None
         try:
             context = zmq.Context(1)
             # Socket facing clients
             frontend = context.socket(zmq.SUB)
-            if in_port is None:
-                in_port = frontend.bind_to_random_port("tcp://*")
-            else:
-                frontend.bind("tcp://*:%d" % in_port)
+            if protocol is Protocols.TCP:
+                if in_port is None:
+                    in_port = frontend.bind_to_random_port(f"{protocol.value}://*")
+                else:
+                    frontend.bind(f"{protocol.value}://*:{in_port}")
+            elif protocol is Protocols.IPC:
+                if in_port is None:
+                    in_port = f"/tmp/{uuid.uuid4()}"
+                frontend.bind(f"{protocol.value}://{in_port}")
 
             frontend.setsockopt_string(zmq.SUBSCRIBE, "")
 
             # Socket facing services
             backend = context.socket(zmq.PUB)
-            if out_port is None:
-                out_port = backend.bind_to_random_port("tcp://*")
-            else:
-                backend.bind("tcp://*:%d" % out_port)
+            if protocol is Protocols.TCP:
+                if out_port is None:
+                    out_port = backend.bind_to_random_port(f"{protocol.value}://*")
+                else:
+                    backend.bind(f"{protocol.value}://*:{out_port}")
+            elif protocol is Protocols.IPC:
+                if out_port is None:
+                    out_port = f"/tmp/{uuid.uuid4()}"
+                frontend.bind(f"{protocol.value}://{out_port}")
+
         except BaseException:
             # Clean up whichever components we have defined so far.
             try:
@@ -247,6 +271,7 @@ def __init__(
         zmq_asyncio=None,
         deserializer=pickle.loads,
         strict=False,
+        protocol=Protocols.TCP,
     ):
         if isinstance(prefix, str):
             raise ValueError("prefix must be bytes, not string")
@@ -257,10 +282,12 @@ def __init__(
             import zmq
         if zmq_asyncio is None:
             import zmq.asyncio as zmq_asyncio
-        if isinstance(address, str):
-            address = address.split(":", maxsplit=1)
+        if isinstance(address, tuple):
+            url = f"{protocol.value}://{address[0]}:{address[1]}"
+        else:
+            url = f"{protocol.value}://address"
         self._deserializer = deserializer
-        self.address = (address[0], int(address[1]))
+        self.address = url
 
         if loop is None:
             loop = asyncio.new_event_loop()
@@ -274,8 +301,7 @@ def __finish_setup():
             self._context = zmq_asyncio.Context()
             self._socket = self._context.socket(zmq.SUB)
 
-            url = "tcp://%s:%d" % self.address
-            self._socket.connect(url)
+            self._socket.connect(self.address)
             self._socket.setsockopt_string(zmq.SUBSCRIBE, "")
 
         self.__factory = __finish_setup