Skip to content

Move the ServiceRegistry into its own module #549

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bdraco merged 1 commit into from
Jun 13, 2021
Merged

Move the ServiceRegistry into its own module #549

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion zeroconf/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@
_TYPE_TXT,
_UNREGISTER_TIME,
)
from .core import NotifyListener, ServiceRegistry, Zeroconf # noqa # import needed for backwards compat
from .core import NotifyListener, Zeroconf # noqa # import needed for backwards compat
from .dns import ( # noqa # import needed for backwards compat
DNSAddress,
DNSCache,
Expand Down Expand Up @@ -110,6 +110,7 @@
ServiceInfo,
ServiceStateChange,
)
from .services.registry import ServiceRegistry # noqa # import needed for backwards compat
from .utils.name import service_type_name # noqa # import needed for backwards compat
from .utils.net import ( # noqa # import needed for backwards compat
add_multicast_member,
Expand Down
93 changes: 2 additions & 91 deletions zeroconf/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,9 +52,10 @@
_UNREGISTER_TIME,
)
from .dns import DNSAddress, DNSCache, DNSIncoming, DNSOutgoing, DNSPointer, DNSQuestion, DNSRecord
from .exceptions import NonUniqueNameException, ServiceNameAlreadyRegistered
from .exceptions import NonUniqueNameException
from .logger import QuietLogger, log
from .services import RecordUpdateListener, ServiceBrowser, ServiceInfo, instance_name_from_service_info
from .services.registry import ServiceRegistry
from .utils.name import service_type_name
from .utils.net import (
IPVersion,
Expand Down Expand Up @@ -226,96 +227,6 @@ def handle_read(self, socket_: socket.socket) -> None:
self.zc.handle_response(msg)


class ServiceRegistry:
"""A registry to keep track of services.

This class exists to ensure services can
be safely added and removed with thread
safety.
"""

def __init__(
self,
) -> None:
"""Create the ServiceRegistry class."""
self.services = {} # type: Dict[str, ServiceInfo]
self.types = {} # type: Dict[str, List]
self.servers = {} # type: Dict[str, List]
self._lock = threading.Lock() # add and remove services thread safe

def add(self, info: ServiceInfo) -> None:
"""Add a new service to the registry."""

with self._lock:
self._add(info)

def remove(self, info: ServiceInfo) -> None:
"""Remove a new service from the registry."""

with self._lock:
self._remove(info)

def update(self, info: ServiceInfo) -> None:
"""Update new service in the registry."""

with self._lock:
self._remove(info)
self._add(info)

def get_service_infos(self) -> List[ServiceInfo]:
"""Return all ServiceInfo."""
return list(self.services.values())

def get_info_name(self, name: str) -> Optional[ServiceInfo]:
"""Return all ServiceInfo for the name."""
return self.services.get(name)

def get_types(self) -> List[str]:
"""Return all types."""
return list(self.types.keys())

def get_infos_type(self, type_: str) -> List[ServiceInfo]:
"""Return all ServiceInfo matching type."""
return self._get_by_index("types", type_)

def get_infos_server(self, server: str) -> List[ServiceInfo]:
"""Return all ServiceInfo matching server."""
return self._get_by_index("servers", server)

def _get_by_index(self, attr: str, key: str) -> List[ServiceInfo]:
"""Return all ServiceInfo matching the index."""
service_infos = []

for name in getattr(self, attr).get(key, [])[:]:
info = self.services.get(name)
# Since we do not get under a lock since it would be
# a performance issue, its possible
# the service can be unregistered during the get
# so we must check if info is None
if info is not None:
service_infos.append(info)

return service_infos

def _add(self, info: ServiceInfo) -> None:
"""Add a new service under the lock."""
lower_name = info.name.lower()
if lower_name in self.services:
raise ServiceNameAlreadyRegistered

self.services[lower_name] = info
self.types.setdefault(info.type, []).append(lower_name)
self.servers.setdefault(info.server, []).append(lower_name)

def _remove(self, info: ServiceInfo) -> None:
"""Remove a service under the lock."""
lower_name = info.name.lower()
old_service_info = self.services[lower_name]
self.types[old_service_info.type].remove(lower_name)
self.servers[old_service_info.server].remove(lower_name)
del self.services[lower_name]


class QueryHandler:
"""Query the ServiceRegistry."""

Expand Down
16 changes: 8 additions & 8 deletions zeroconf/services.py → zeroconf/services/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
from collections import OrderedDict
from typing import Any, Callable, Dict, List, Optional, Set, TYPE_CHECKING, Tuple, Union, cast

from .const import (
from ..const import (
_BROWSER_BACKOFF_LIMIT,
_BROWSER_TIME,
_CLASS_IN,
Expand All @@ -46,20 +46,20 @@
_TYPE_SRV,
_TYPE_TXT,
)
from .dns import DNSAddress, DNSOutgoing, DNSPointer, DNSQuestion, DNSRecord, DNSService, DNSText
from .exceptions import BadTypeInNameException
from .utils.name import service_type_name
from .utils.net import (
from ..dns import DNSAddress, DNSOutgoing, DNSPointer, DNSQuestion, DNSRecord, DNSService, DNSText
from ..exceptions import BadTypeInNameException
from ..utils.name import service_type_name
from ..utils.net import (
IPVersion,
_encode_address,
_is_v6_address,
)
from .utils.struct import int2byte
from .utils.time import current_time_millis, millis_to_seconds
from ..utils.struct import int2byte
from ..utils.time import current_time_millis, millis_to_seconds

if TYPE_CHECKING:
# https://github.com/PyCQA/pylint/issues/3525
from . import ( # pylint: disable=cyclic-import
from .. import ( # pylint: disable=cyclic-import
ServiceListener,
Zeroconf,
)
Expand Down
118 changes: 118 additions & 0 deletions zeroconf/services/registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
""" Multicast DNS Service Discovery for Python, v0.14-wmcbrine
Copyright 2003 Paul Scott-Murphy, 2014 William McBrine

This module provides a framework for the use of DNS Service Discovery
using IP multicast.

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301
USA
"""

import threading
from typing import Dict, List, Optional


from ..exceptions import ServiceNameAlreadyRegistered
from ..services import ServiceInfo


class ServiceRegistry:
"""A registry to keep track of services.

This class exists to ensure services can
be safely added and removed with thread
safety.
"""

def __init__(
self,
) -> None:
"""Create the ServiceRegistry class."""
self.services = {} # type: Dict[str, ServiceInfo]
self.types = {} # type: Dict[str, List]
self.servers = {} # type: Dict[str, List]
self._lock = threading.Lock() # add and remove services thread safe

def add(self, info: ServiceInfo) -> None:
"""Add a new service to the registry."""

with self._lock:
self._add(info)

def remove(self, info: ServiceInfo) -> None:
"""Remove a new service from the registry."""

with self._lock:
self._remove(info)

def update(self, info: ServiceInfo) -> None:
"""Update new service in the registry."""

with self._lock:
self._remove(info)
self._add(info)

def get_service_infos(self) -> List[ServiceInfo]:
"""Return all ServiceInfo."""
return list(self.services.values())

def get_info_name(self, name: str) -> Optional[ServiceInfo]:
"""Return all ServiceInfo for the name."""
return self.services.get(name)

def get_types(self) -> List[str]:
"""Return all types."""
return list(self.types.keys())

def get_infos_type(self, type_: str) -> List[ServiceInfo]:
"""Return all ServiceInfo matching type."""
return self._get_by_index("types", type_)

def get_infos_server(self, server: str) -> List[ServiceInfo]:
"""Return all ServiceInfo matching server."""
return self._get_by_index("servers", server)

def _get_by_index(self, attr: str, key: str) -> List[ServiceInfo]:
"""Return all ServiceInfo matching the index."""
service_infos = []

for name in getattr(self, attr).get(key, [])[:]:
info = self.services.get(name)
# Since we do not get under a lock since it would be
# a performance issue, its possible
# the service can be unregistered during the get
# so we must check if info is None
if info is not None:
service_infos.append(info)

return service_infos

def _add(self, info: ServiceInfo) -> None:
"""Add a new service under the lock."""
lower_name = info.name.lower()
if lower_name in self.services:
raise ServiceNameAlreadyRegistered

self.services[lower_name] = info
self.types.setdefault(info.type, []).append(lower_name)
self.servers.setdefault(info.server, []).append(lower_name)

def _remove(self, info: ServiceInfo) -> None:
"""Remove a service under the lock."""
lower_name = info.name.lower()
old_service_info = self.services[lower_name]
self.types[old_service_info.type].remove(lower_name)
self.servers[old_service_info.server].remove(lower_name)
del self.services[lower_name]