Break apart Zeroconf.handle_query to reduce branching (#462)

bdraco · web-flow · commit c1ed987ede34 · 2021-06-05T16:08:27.000-10:00
diff --git a/zeroconf/__init__.py b/zeroconf/__init__.py
@@ -2948,9 +2948,59 @@ def handle_response(self, msg: DNSIncoming) -> None:  # pylint: disable=too-many
             for record in removes:
                 self.cache.remove(record)
 
-    def handle_query(  # pylint: disable=too-many-branches
-        self, msg: DNSIncoming, addr: Optional[str], port: int
-    ) -> None:
+    def _answer_service_type_enumeration_query(self, msg: DNSIncoming, out: DNSOutgoing) -> None:
+        """Provide an answer to a service type enumeration query.
+
+        https://datatracker.ietf.org/doc/html/rfc6763#section-9
+        """
+        for stype in self.registry.get_types():
+            out.add_answer(
+                msg,
+                DNSPointer(
+                    _SERVICE_TYPE_ENUMERATION_NAME,
+                    _TYPE_PTR,
+                    _CLASS_IN,
+                    _DNS_OTHER_TTL,
+                    stype,
+                ),
+            )
+
+    def _answer_ptr_query(self, msg: DNSIncoming, out: DNSOutgoing, question: DNSQuestion) -> None:
+        """Answer a PTR query."""
+        for service in self.registry.get_infos_type(question.name):
+            out.add_answer(msg, service.dns_pointer())
+            # Add recommended additional answers according to
+            # https://tools.ietf.org/html/rfc6763#section-12.1.
+            out.add_additional_answer(service.dns_service())
+            out.add_additional_answer(service.dns_text())
+            for dns_address in service.dns_addresses():
+                out.add_additional_answer(dns_address)
+
+    def _answer_non_ptr_query(self, msg: DNSIncoming, out: DNSOutgoing, question: DNSQuestion) -> None:
+        """Answer a query any query other then PTR.
+
+        Add answer(s) for A, AAAA, SRV, or TXT queries.
+        """
+        name_to_find = question.name.lower()
+        # Answer A record queries for any service addresses we know
+        if question.type in (_TYPE_A, _TYPE_ANY):
+            for service in self.registry.get_infos_server(name_to_find):
+                for dns_address in service.dns_addresses():
+                    out.add_answer(msg, dns_address)
+
+        service = self.registry.get_info_name(name_to_find)  # type: ignore
+        if service is None:
+            return
+
+        if question.type in (_TYPE_SRV, _TYPE_ANY):
+            out.add_answer(msg, service.dns_service())
+        if question.type in (_TYPE_TXT, _TYPE_ANY):
+            out.add_answer(msg, service.dns_text())
+        if question.type == _TYPE_SRV:
+            for dns_address in service.dns_addresses():
+                out.add_additional_answer(dns_address)
+
+    def handle_query(self, msg: DNSIncoming, addr: Optional[str], port: int) -> None:
         """Deal with incoming query packets.  Provides a response if
         possible."""
         # Support unicast client responses
@@ -2965,48 +3015,12 @@ def handle_query(  # pylint: disable=too-many-branches
         for question in msg.questions:
             if question.type == _TYPE_PTR:
                 if question.name == _SERVICE_TYPE_ENUMERATION_NAME:
-                    for stype in self.registry.get_types():
-                        out.add_answer(
-                            msg,
-                            DNSPointer(
-                                _SERVICE_TYPE_ENUMERATION_NAME,
-                                _TYPE_PTR,
-                                _CLASS_IN,
-                                _DNS_OTHER_TTL,
-                                stype,
-                            ),
-                        )
-                    continue
-
-                for service in self.registry.get_infos_type(question.name):
-                    out.add_answer(msg, service.dns_pointer())
-                    # Add recommended additional answers according to
-                    # https://tools.ietf.org/html/rfc6763#section-12.1.
-                    out.add_additional_answer(service.dns_service())
-                    out.add_additional_answer(service.dns_text())
-                    for dns_address in service.dns_addresses():
-                        out.add_additional_answer(dns_address)
-
-                continue
-
-            name_to_find = question.name.lower()
-            # Answer A record queries for any service addresses we know
-            if question.type in (_TYPE_A, _TYPE_ANY):
-                for service in self.registry.get_infos_server(name_to_find):
-                    for dns_address in service.dns_addresses():
-                        out.add_answer(msg, dns_address)
-
-            service = self.registry.get_info_name(name_to_find)  # type: ignore
-            if service is None:
+                    self._answer_service_type_enumeration_query(msg, out)
+                else:
+                    self._answer_ptr_query(msg, out, question)
                 continue
 
-            if question.type in (_TYPE_SRV, _TYPE_ANY):
-                out.add_answer(msg, service.dns_service())
-            if question.type in (_TYPE_TXT, _TYPE_ANY):
-                out.add_answer(msg, service.dns_text())
-            if question.type == _TYPE_SRV:
-                for dns_address in service.dns_addresses():
-                    out.add_additional_answer(dns_address)
+            self._answer_non_ptr_query(msg, out, question)
 
         if out is not None and out.answers:
             out.id = msg.id
