diff --git a/src/dsf/connections/base_connection.py b/src/dsf/connections/base_connection.py index 38ac463..2cf3f19 100644 --- a/src/dsf/connections/base_connection.py +++ b/src/dsf/connections/base_connection.py @@ -20,12 +20,16 @@ def __init__(self, debug: bool = False, timeout: int = 3): self.id = None self.input = "" - def connect(self, init_message: client_init_messages.ClientInitMessage, socket_file: str): + def connect(self, init_message: client_init_messages.ClientInitMessage, socket_file: str, timeout: int = 0): """Establishes a connection to the given UNIX socket file""" self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) self.socket.connect(socket_file) - self.socket.setblocking(True) + if timeout == 0: + self.socket.setblocking(True) + else: + self.socket.setblocking(False) + self.socket.settimeout(timeout) server_init_msg = server_init_message.ServerInitMessage.from_json( json.loads(self.socket.recv(50).decode("utf8")) ) @@ -75,7 +79,10 @@ def send(self, msg): def receive(self, cls): """Receive a deserialized object from the server""" json_string = self.receive_json() - return cls.from_json(json.loads(json_string)) + try: + return cls.from_json(json.loads(json_string)) + except Exception as e: + return None def receive_response(self): """Receive a base response from the server""" @@ -108,7 +115,7 @@ def receive_json(self) -> str: part = self.socket.recv(BUFF_SIZE) data += part except socket.timeout: - pass + return None except Exception as e: raise e # either 0 or end of data diff --git a/src/dsf/connections/intercept_connection.py b/src/dsf/connections/intercept_connection.py index 6ef0154..d8db898 100644 --- a/src/dsf/connections/intercept_connection.py +++ b/src/dsf/connections/intercept_connection.py @@ -46,7 +46,7 @@ def __init__( self.auto_evaluate_expression = auto_evaluate_expression self.priority_codes = priority_codes - def connect(self, socket_file: str = SOCKET_FILE): # noqa + def connect(self, timeout: int = 0, socket_file: str = SOCKET_FILE): # noqa """Establishes a connection to the given UNIX socket file""" iim = client_init_messages.intercept_init_message( self.interception_mode, @@ -56,7 +56,7 @@ def connect(self, socket_file: str = SOCKET_FILE): # noqa self.auto_flush, self.auto_evaluate_expression ) - return super().connect(iim, socket_file) + return super().connect(iim, socket_file, timeout) def receive_code(self) -> commands.code.Code: """Wait for a code to be intercepted and read it"""