diff --git a/royalnet/herald/link.py b/royalnet/herald/link.py index ce234282..690f5ab5 100644 --- a/royalnet/herald/link.py +++ b/royalnet/herald/link.py @@ -1,5 +1,4 @@ import asyncio -import websockets import uuid import functools import logging as _logging @@ -11,6 +10,11 @@ from .broadcast import Broadcast from .errors import ConnectionClosedError, InvalidServerResponseError from .config import Config +try: + import websockets +except ImportError: + websockets = None + log = _logging.getLogger(__name__) @@ -53,9 +57,11 @@ def requires_identification(func): class Link: def __init__(self, config: Config, request_handler, *, loop: asyncio.AbstractEventLoop = None): + if websockets is None: + raise ImportError("'websockets' extra is not installed") self.config: Config = config self.nid: str = str(uuid.uuid4()) - self.websocket: typing.Optional[websockets.WebSocketClientProtocol] = None + self.websocket: typing.Optional["websockets.WebSocketClientProtocol"] = None self.request_handler: typing.Callable[[typing.Union[Request, Broadcast]], typing.Awaitable[Response]] = request_handler self._pending_requests: typing.Dict[str, PendingRequest] = {} diff --git a/royalnet/herald/server.py b/royalnet/herald/server.py index dff0bd29..d43cbe8b 100644 --- a/royalnet/herald/server.py +++ b/royalnet/herald/server.py @@ -1,5 +1,4 @@ import typing -import websockets import re import datetime import uuid @@ -8,14 +7,19 @@ import logging as _logging from .package import Package from .config import Config +try: + import websockets +except ImportError: + websockets = None + log = _logging.getLogger(__name__) class ConnectedClient: """The :py:class:`Server`-side representation of a connected :py:class:`Link`.""" - def __init__(self, socket: websockets.WebSocketServerProtocol): - self.socket: websockets.WebSocketServerProtocol = socket + def __init__(self, socket: "websockets.WebSocketServerProtocol"): + self.socket: "websockets.WebSocketServerProtocol" = socket self.nid: typing.Optional[str] = None self.link_type: typing.Optional[str] = None self.connection_datetime: datetime.datetime = datetime.datetime.now() @@ -57,7 +61,7 @@ class Server: matching = [client for client in self.identified_clients if client.link_type == link_type] return matching or [] - async def listener(self, websocket: websockets.server.WebSocketServerProtocol, path): + async def listener(self, websocket: "websockets.server.WebSocketServerProtocol", path): log.info(f"{websocket.remote_address} connected to the server.") connected_client = ConnectedClient(websocket) # Wait for identification