python-zeroconf · bdraco · Sep 3, 2023 · Sep 3, 2023 · Sep 3, 2023
diff --git a/build_ext.py b/build_ext.py
@@ -30,6 +30,7 @@ def build(setup_kwargs: Any) -> None:
                         "src/zeroconf/_protocol/incoming.py",
                         "src/zeroconf/_protocol/outgoing.py",
                         "src/zeroconf/_handlers/record_manager.py",
+                        "src/zeroconf/_handlers/query_handler.py",
                         "src/zeroconf/_services/registry.py",
                         "src/zeroconf/_utils/time.py",
                     ],

diff --git a/src/zeroconf/_handlers/query_handler.pxd b/src/zeroconf/_handlers/query_handler.pxd
@@ -0,0 +1,64 @@
+
+import cython
+
+from .._cache cimport DNSCache
+from .._dns cimport DNSPointer, DNSQuestion, DNSRecord, DNSRRSet
+from .._history cimport QuestionHistory
+from .._protocol.incoming cimport DNSIncoming
+from .._services.registry cimport ServiceRegistry
+
+
+cdef object TYPE_CHECKING, QuestionAnswers
+cdef cython.uint _ONE_SECOND, _TYPE_PTR, _TYPE_ANY, _TYPE_A, _TYPE_AAAA, _TYPE_SRV, _TYPE_TXT
+cdef str _SERVICE_TYPE_ENUMERATION_NAME
+cdef cython.set _RESPOND_IMMEDIATE_TYPES
+
+cdef class _QueryResponse:
+
+    cdef object _is_probe
+    cdef DNSIncoming _msg
+    cdef float _now
+    cdef DNSCache _cache
+    cdef cython.dict _additionals
+    cdef cython.set _ucast
+    cdef cython.set _mcast_now
+    cdef cython.set _mcast_aggregate
+    cdef cython.set _mcast_aggregate_last_second
+
+    cpdef add_qu_question_response(self, cython.dict answers)
+
+    cpdef add_ucast_question_response(self, cython.dict answers)
+
+    cpdef add_mcast_question_response(self, cython.dict answers)
+
+    @cython.locals(maybe_entry=DNSRecord)
+    cpdef _has_mcast_within_one_quarter_ttl(self, DNSRecord record)
+
+    @cython.locals(maybe_entry=DNSRecord)
+    cpdef _has_mcast_record_in_last_second(self, DNSRecord record)
+
+    cpdef answers(self)
+
+cdef class QueryHandler:
+
+    cdef ServiceRegistry registry
+    cdef DNSCache cache
+    cdef QuestionHistory question_history
+
+    cdef _add_service_type_enumeration_query_answers(self, cython.dict answer_set, DNSRRSet known_answers)
+
+    cdef _add_pointer_answers(self, str lower_name, cython.dict answer_set, DNSRRSet known_answers)
+
+    cdef _add_address_answers(self, str lower_name, cython.dict answer_set, DNSRRSet known_answers, cython.uint type_)
+
+    @cython.locals(question_lower_name=str, type_=cython.uint)
+    cdef _answer_question(self, DNSQuestion question, DNSRRSet known_answers)
+
+    @cython.locals(
+        msg=DNSIncoming,
+        question=DNSQuestion,
+        answer_set=cython.dict,
+        known_answers=DNSRRSet,
+        known_answers_set=cython.set,
+    )
+    cpdef async_response(self, cython.list msgs, object unicast_source)
diff --git a/src/zeroconf/_handlers/query_handler.py b/src/zeroconf/_handlers/query_handler.py
@@ -46,6 +46,8 @@
 
 _RESPOND_IMMEDIATE_TYPES = {_TYPE_NSEC, _TYPE_SRV, *_ADDRESS_RECORD_TYPES}
 
+_int = int
+
 
 class _QueryResponse:
     """A pair for unicast and multicast DNSOutgoing responses."""
@@ -113,17 +115,11 @@ def answers(
         self,
     ) -> QuestionAnswers:
         """Return answer sets that will be queued."""
-        return QuestionAnswers(
-            *(
-                {record: self._additionals[record] for record in rrset}
-                for rrset in (
-                    self._ucast,
-                    self._mcast_now,
-                    self._mcast_aggregate,
-                    self._mcast_aggregate_last_second,
-                )
-            )
-        )
+        ucast = {r: self._additionals[r] for r in self._ucast}
+        mcast_now = {r: self._additionals[r] for r in self._mcast_now}
+        mcast_aggregate = {r: self._additionals[r] for r in self._mcast_aggregate}
+        mcast_aggregate_last_second = {r: self._additionals[r] for r in self._mcast_aggregate_last_second}
+        return QuestionAnswers(ucast, mcast_now, mcast_aggregate, mcast_aggregate_last_second)
 
     def _has_mcast_within_one_quarter_ttl(self, record: DNSRecord) -> bool:
         """Check to see if a record has been mcasted recently.
@@ -197,7 +193,7 @@ def _add_address_answers(
         lower_name: str,
         answer_set: _AnswerWithAdditionalsType,
         known_answers: DNSRRSet,
-        type_: int,
+        type_: _int,
     ) -> None:
         """Answer A/AAAA/ANY question."""
         for service in self.registry.async_get_infos_server(lower_name):

diff --git a/src/zeroconf/_history.pxd b/src/zeroconf/_history.pxd
@@ -1,16 +1,18 @@
 import cython
 
+from ._dns cimport DNSQuestion
+
 
 cdef cython.double _DUPLICATE_QUESTION_INTERVAL
 
 cdef class QuestionHistory:
 
     cdef cython.dict _history
 
+    cpdef add_question_at_time(self, DNSQuestion question, float now, cython.set known_answers)
 
     @cython.locals(than=cython.double, previous_question=cython.tuple, previous_known_answers=cython.set)
-    cpdef suppresses(self, object question, cython.double now, cython.set known_answers)
-
+    cpdef suppresses(self, DNSQuestion question, cython.double now, cython.set known_answers)
 
     @cython.locals(than=cython.double, now_known_answers=cython.tuple)
     cpdef async_expire(self, cython.double now)
diff --git a/src/zeroconf/_services/registry.pxd b/src/zeroconf/_services/registry.pxd
@@ -16,3 +16,11 @@ cdef class ServiceRegistry:
     cdef _add(self, object info)
 
     cdef _remove(self, cython.list infos)
+
+    cpdef async_get_info_name(self, str name)
+
+    cpdef async_get_types(self)
+
+    cpdef async_get_infos_type(self, str type_)
+
+    cpdef async_get_infos_server(self, str server)