From 47327cf693042c126b5d8f026bd1568f69716fb0 Mon Sep 17 00:00:00 2001
From: Stefano Pigozzi <ste.pigozzi@gmail.com>
Date: Thu, 14 Nov 2019 12:50:44 +0100
Subject: [PATCH] #105: Add optional imports to herald

---
 royalnet/herald/link.py   | 10 ++++++++--
 royalnet/herald/server.py | 12 ++++++++----
 2 files changed, 16 insertions(+), 6 deletions(-)

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