From 210cefc1a40977281db2a9682aa67f81e5f94cc7 Mon Sep 17 00:00:00 2001
From: Fabio Manganiello <info@fabiomanganiello.com>
Date: Wed, 3 Mar 2021 19:19:56 +0100
Subject: [PATCH] General improvements on the Zeroconf plugin and backend

---
 platypush/backend/__init__.py | 31 +++++++++++++++++++++++++++----
 platypush/plugins/zeroconf.py | 15 +++++++++------
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/platypush/backend/__init__.py b/platypush/backend/__init__.py
index a779cfc0..fd5d9cab 100644
--- a/platypush/backend/__init__.py
+++ b/platypush/backend/__init__.py
@@ -10,7 +10,7 @@ import threading
 import time
 
 from threading import Thread
-from typing import Optional
+from typing import Optional, Dict
 
 from platypush.bus import Bus
 from platypush.config import Config
@@ -313,9 +313,31 @@ class Backend(Thread, EventGenerator):
         s.close()
         return addr
 
-    def register_service(self, port: Optional[int] = None, name: Optional[str] = None, udp: bool = False):
+    def register_service(self,
+                         port: Optional[int] = None,
+                         name: Optional[str] = None,
+                         srv_type: Optional[str] = None,
+                         srv_name: Optional[str] = None,
+                         udp: bool = False,
+                         properties: Optional[Dict] = None):
         """
         Initialize the Zeroconf service configuration for this backend.
+
+        :param port: Service listen port (default: the backend ``port`` attribute if available, or ``None``).
+        :param name: Service short name (default: backend name).
+        :param srv_type: Service type (default: ``_platypush-{name}._{proto}.local.``).
+        :param srv_name: Full service name (default: ``{hostname or device_id}.{type}``).
+        :param udp: Set to True if this is a UDP service.
+        :param properties: Extra properties to be passed on the service. Default:
+
+            .. code-block:: json
+
+                {
+                    "name": "Platypush",
+                    "vendor": "Platypush",
+                    "version": "{platypush_version}"
+                }
+
         """
         try:
             from zeroconf import ServiceInfo, Zeroconf
@@ -329,11 +351,12 @@ class Backend(Thread, EventGenerator):
             'name': 'Platypush',
             'vendor': 'Platypush',
             'version': __version__,
+            **(properties or {}),
         }
 
         name = name or re.sub(r'Backend$', '', self.__class__.__name__).lower()
-        srv_type = '_platypush-{name}._{proto}.local.'.format(name=name, proto='udp' if udp else 'tcp')
-        srv_name = '{host}.{type}'.format(host=self.device_id, type=srv_type)
+        srv_type = srv_type or '_platypush-{name}._{proto}.local.'.format(name=name, proto='udp' if udp else 'tcp')
+        srv_name = srv_name or '{host}.{type}'.format(host=self.device_id, type=srv_type)
 
         if port:
             srv_port = port
diff --git a/platypush/plugins/zeroconf.py b/platypush/plugins/zeroconf.py
index ba85d3e9..e386b355 100644
--- a/platypush/plugins/zeroconf.py
+++ b/platypush/plugins/zeroconf.py
@@ -1,14 +1,14 @@
 import queue
 import socket
 import time
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Any, Optional, Union
 
 import zeroconf
 from zeroconf import Zeroconf, ServiceInfo, ServiceBrowser
 
 from platypush.context import get_bus
 from platypush.message.event.zeroconf import ZeroconfServiceAddedEvent, ZeroconfServiceRemovedEvent, \
-    ZeroconfServiceUpdatedEvent, ZeroconfEvent
+    ZeroconfServiceUpdatedEvent
 from platypush.plugins import Plugin, action
 
 
@@ -82,11 +82,11 @@ class ZeroconfPlugin(Plugin):
         return list(zeroconf.ZeroconfServiceTypes.find(timeout=timeout))
 
     @action
-    def discover_service(self, service: str, timeout: Optional[int] = 5) -> Dict[str, Any]:
+    def discover_service(self, service: Union[str, list], timeout: Optional[int] = 5) -> Dict[str, Any]:
         """
         Find all the services matching the specified type.
 
-        :param service: Service type (e.g. ``_http._tcp.local.``).
+        :param service: Service type (e.g. ``_http._tcp.local.``) or list of service types.
         :param timeout: Browser timeout in seconds (default: 5). Specify None for no timeout - in such case the
             discovery will loop forever and generate events upon service changes.
         :return: A ``service_type -> [service_names]`` mapping. Example:
@@ -121,11 +121,13 @@ class ZeroconfPlugin(Plugin):
         evt_queue = queue.Queue()
         zc = Zeroconf()
         listener = ZeroconfListener(evt_queue=evt_queue)
-        browser = ServiceBrowser(zc, service, listener)
         discovery_start = time.time()
         services = {}
+        browser = None
 
         try:
+            browser = ServiceBrowser(zc, service, listener)
+
             while timeout and time.time() - discovery_start < timeout:
                 to = discovery_start + timeout - time.time() if timeout else None
                 try:
@@ -145,7 +147,8 @@ class ZeroconfPlugin(Plugin):
                     if not services:
                         self.logger.warning('No such service discovered: {}'.format(service))
         finally:
-            browser.cancel()
+            if browser:
+                browser.cancel()
             zc.close()
             self._discovery_in_progress = False