diff --git a/truenas_api_client/__init__.py b/truenas_api_client/__init__.py index b84923c..c7cc3a1 100644 --- a/truenas_api_client/__init__.py +++ b/truenas_api_client/__init__.py @@ -71,7 +71,8 @@ class Client: """Implicit wrapper of either a `JSONRPCClient` or a `LegacyClient`.""" def __init__(self, uri: str | None = None, reserved_ports=False, private_methods=False, py_exceptions=False, - log_py_exceptions=False, call_timeout: float | UndefinedType = undefined, verify_ssl=True): + log_py_exceptions=False, call_timeout: float | UndefinedType = undefined, verify_ssl=True, + ping_interval: float | int = 0, reconnect: int = None): """Initialize either a `JSONRPCClient` or a `LegacyClient`. Use `JSONRPCClient` unless `uri` ends with '/websocket'. @@ -86,6 +87,8 @@ def __init__(self, uri: str | None = None, reserved_ports=False, private_methods call_timeout: Number of seconds to allow an API call before timing out. Can be overridden on a per-call basis. Defaults to `CALL_TIMEOUT`. verify_ssl: `True` if SSL certificate should be verified before connecting. + ping_interval: Number of seconds between WebSocket ping frames. Defaults to 0 (disabled). + reconnect: Attempt to reconnect every n seconds. Defaults to None (no reconnects). Raises: ClientException: `WSClient` timed out or some other connection error occurred. @@ -99,7 +102,7 @@ def __init__(self, uri: str | None = None, reserved_ports=False, private_methods self.uri_check(uri, py_exceptions) self.__client = client_class(uri, reserved_ports, private_methods, py_exceptions, log_py_exceptions, - call_timeout, verify_ssl) + call_timeout, verify_ssl, ping_interval, reconnect) def uri_check(self, uri: str | None, py_exceptions: bool): # We pickle_load when handling py_exceptions, reduce risk of MITM on client causing a pickle.load @@ -123,7 +126,8 @@ class WSClient: The object used by `JSONRPCClient` to send and receive data. """ - def __init__(self, url: str, *, client: 'JSONRPCClient', reserved_ports: bool = False, verify_ssl: bool = True): + def __init__(self, url: str, *, client: 'JSONRPCClient', reserved_ports: bool = False, verify_ssl: bool = True, + ping_interval: float | int = 0, reconnect: int = None): """Initialize a `WSClient`. Args: @@ -131,12 +135,16 @@ def __init__(self, url: str, *, client: 'JSONRPCClient', reserved_ports: bool = client: Reference to the `JSONRPCClient` instance that uses this object. reserved_ports: `True` if the `socket` should bind to a reserved port, i.e. 600-1024. verify_ssl: `True` if SSL certificate should be verified before connecting. + ping_interval: Number of seconds between WebSocket ping frames. Defaults to 0 (disabled). + reconnect: Attempt to reconnect every n seconds. Defaults to None (no reconnects). """ self.url = url self.client = client self.reserved_ports = reserved_ports self.verify_ssl = verify_ssl + self.ping_interval = ping_interval + self.reconnect = reconnect self.socket: socket.socket self.app: WebSocketApp @@ -177,7 +185,7 @@ def connect(self): on_error=self._on_error, on_close=self._on_close, ) - Thread(daemon=True, target=self.app.run_forever).start() + Thread(daemon=True, target=self.app.run_forever, kwargs={"ping_interval": self.ping_interval, "reconnect": self.reconnect}).start() def send(self, data: bytes | str): """Send data to the server by calling `WebSocketApp.send()`. @@ -398,7 +406,8 @@ class JSONRPCClient: """ def __init__(self, uri: str | None = None, reserved_ports=False, private_methods=False, py_exceptions=False, - log_py_exceptions=False, call_timeout: float | UndefinedType = undefined, verify_ssl=True): + log_py_exceptions=False, call_timeout: float | UndefinedType = undefined, verify_ssl=True, + ping_interval: float | int = 0, reconnect: int = None): """Initialize a `JSONRPCClient`. Args: @@ -411,6 +420,8 @@ def __init__(self, uri: str | None = None, reserved_ports=False, private_methods call_timeout: Number of seconds to allow an API call before timing out. Can be overridden on a per-call basis. Defaults to `CALL_TIMEOUT`. verify_ssl: `True` if SSL certificate should be verified before connecting. + ping_interval: Number of seconds between WebSocket ping frames. Defaults to 0 (disabled). + reconnect: Attempt to reconnect every n seconds. Defaults to None (no reconnects). Raises: ClientException: `WSClient` timed out or some other connection error occurred. @@ -442,6 +453,8 @@ def __init__(self, uri: str | None = None, reserved_ports=False, private_methods client=self, reserved_ports=reserved_ports, verify_ssl=verify_ssl, + ping_interval=ping_interval, + reconnect=reconnect, ) self._ws.connect() self._connected.wait(10) diff --git a/truenas_api_client/legacy.py b/truenas_api_client/legacy.py index 8657c40..ff3c4fd 100644 --- a/truenas_api_client/legacy.py +++ b/truenas_api_client/legacy.py @@ -28,11 +28,13 @@ class WSClient: - def __init__(self, url, *, client, reserved_ports=False, verify_ssl=True): + def __init__(self, url, *, client, reserved_ports=False, verify_ssl=True, ping_interval: float | int = 0, reconnect: int = None): self.url = url self.client = client self.reserved_ports = reserved_ports self.verify_ssl = verify_ssl + self.ping_interval = ping_interval + self.reconnect = reconnect self.socket = None self.app = None @@ -68,7 +70,7 @@ def connect(self): on_error=self._on_error, on_close=self._on_close, ) - Thread(daemon=True, target=self.app.run_forever).start() + Thread(daemon=True, target=self.app.run_forever, kwargs={"ping_interval": self.ping_interval, "reconnect": self.reconnect}).start() def send(self, data): return self.app.send(data) @@ -179,7 +181,8 @@ def result(self): class LegacyClient: def __init__(self, uri=None, reserved_ports=False, private_methods=False, py_exceptions=False, - log_py_exceptions=False, call_timeout: float | UndefinedType = undefined, verify_ssl=True): + log_py_exceptions=False, call_timeout: float | UndefinedType = undefined, verify_ssl=True, + ping_interval: float | int = 0, reconnect: int = None): """ Arguments: :reserved_ports(bool): should the local socket used a reserved port @@ -207,6 +210,8 @@ def __init__(self, uri=None, reserved_ports=False, private_methods=False, py_exc client=self, reserved_ports=reserved_ports, verify_ssl=verify_ssl, + ping_interval=ping_interval, + reconnect=reconnect, ) self._ws.connect() self._connected.wait(10)