pyzmq 25 compatibility

Handle a new keyword argument to `Context.socket` in our subclass.

Author: chrisjbillington

labscript_utils/ls_zprocess.py

@@ -18,7 +18,7 @@
 
 import zprocess
 import zprocess.process_tree
-from zprocess.security import SecureContext
+from zprocess.security import SecureContext, SecureSocket
 from labscript_utils.labconfig import LabConfig
 from labscript_utils import dedent
 import zprocess.zlog
@@ -248,10 +248,19 @@ def instance(cls):
         # Super required to call unbound class method of parent class:
         return super(Context, cls).instance(shared_secret=config['shared_secret'])
 
-    def socket(self, *args, **kwargs):
-        config = get_config()
-        kwargs['allow_insecure'] = config['allow_insecure']
-        return SecureContext.socket(self, *args, **kwargs)
+    def socket(self, socket_type, socket_class=None, **kwargs):
+        # socket_class kwarg introduced in pyzmq 25. Pass it through if it was given,
+        # otherwise don't.
+        if socket_class is not None:
+            kwargs['socket_class'] = socket_class
+        # Only insert our security-related args to the socket if we know it's going to
+        # be a SecureSocket. If caller has explicitly requested a different socket type
+        # (e.g since pyzmq 25, ThreadAuthenticator sets up an internal socket by calling
+        # `Context.socket(..., socket_class=zmq.Socket)), then don't.`
+        if socket_class is None or issubclass(socket_class, SecureContext):
+            config = get_config()
+            kwargs['allow_insecure'] = config['allow_insecure']
+        return SecureContext.socket(self, socket_type=socket_type, **kwargs)
 
 
 def Lock(*args, **kwargs):