This repository has been archived on 2023-03-26. You can view files and clone it, but cannot push or open issues or pull requests.
ipc-unix/ipc_unix/pubsub.py
2018-12-07 23:18:29 +00:00

66 lines
1.8 KiB
Python

import select
import socket
import ujson
from ipc_unix.utils import read_payload
class Subscriber:
def __init__(self, socket_path):
self.socket_path = socket_path
self.socket = socket.socket(
socket.AF_UNIX, type=socket.SOCK_STREAM | socket.SOCK_NONBLOCK
)
self.socket.connect(self.socket_path)
def listen(self):
while True:
yield self.get_message()
def get_message(self):
return read_payload(self.socket)
def close(self):
self.socket.close()
class Publisher:
def __init__(self, socket_path):
self.socket_path = socket_path
self.master_socket = socket.socket(
socket.AF_UNIX, type=socket.SOCK_STREAM | socket.SOCK_NONBLOCK
)
self.master_socket.bind(self.socket_path)
self.master_socket.listen()
self.connections = []
def close(self):
self.master_socket.close()
self.connections.clear()
def accept_new_connection(self):
readable, _, _ = select.select([self.master_socket], [], [], 1)
if self.master_socket in readable:
new_socket, _ = self.master_socket.accept()
self.connections.append(new_socket)
def write(self, message):
self.accept_new_connection()
_, writable, errorable = select.select([], self.connections, [], 1)
dead_sockets = []
if writable:
data = ujson.dumps(message).encode() + b"\n"
for sock in writable:
try:
sock.send(data)
except BrokenPipeError:
dead_sockets.append(sock)
for sock in dead_sockets:
if sock in self.connections:
self.connections.remove(sock)
sock.close()