Improve get_rrset/find_rrset API. (#922)

* Improve get_rrset/find_rrset API.

This allows most of the parameters to be specified as strings, matching
the interface for dns.message.make_query().

* Remove unneeded "the_section".

There's no need to use a separate internal variable for the section;
mypy doesn't complain about reuse.

diff --git a/dns/message.py b/dns/message.py
index 082fb7b6dbe188975e682bd202227083f1a35b84..3a6f4273f6a2c7e4693c3a91c1ee9fd06f509fc7 100644 (file)
--- a/dns/message.py
+++ b/dns/message.py
@@ -135,7 +135,7 @@ IndexKeyType = Tuple[
     Optional[dns.rdataclass.RdataClass],
 ]
 IndexType = Dict[IndexKeyType, dns.rrset.RRset]
-SectionType = Union[int, List[dns.rrset.RRset]]
+SectionType = Union[int, str, List[dns.rrset.RRset]]
 
 
 class Message:
@@ -348,27 +348,29 @@ class Message:
         deleting: Optional[dns.rdataclass.RdataClass] = None,
         create: bool = False,
         force_unique: bool = False,
+        idna_codec: Optional[dns.name.IDNACodec] = None,
     ) -> dns.rrset.RRset:
         """Find the RRset with the given attributes in the specified section.
 
-        *section*, an ``int`` section number, or one of the section
-        attributes of this message.  This specifies the
+        *section*, an ``int`` section number, a ``str`` section name, or one of
+        the section attributes of this message.  This specifies the
         the section of the message to search.  For example::
 
             my_message.find_rrset(my_message.answer, name, rdclass, rdtype)
             my_message.find_rrset(dns.message.ANSWER, name, rdclass, rdtype)
+            my_message.find_rrset("ANSWER", name, rdclass, rdtype)
 
-        *name*, a ``dns.name.Name``, the name of the RRset.
+        *name*, a ``dns.name.Name`` or ``str``, the name of the RRset.
 
-        *rdclass*, an ``int``, the class of the RRset.
+        *rdclass*, an ``int`` or ``str``, the class of the RRset.
 
-        *rdtype*, an ``int``, the type of the RRset.
+        *rdtype*, an ``int`` or ``str``, the type of the RRset.
 
-        *covers*, an ``int`` or ``None``, the covers value of the RRset.
-        The default is ``None``.
+        *covers*, an ``int`` or ``str``, the covers value of the RRset.
+        The default is ``dns.rdatatype.NONE``.
 
-        *deleting*, an ``int`` or ``None``, the deleting value of the RRset.
-        The default is ``None``.
+        *deleting*, an ``int``, ``str``, or ``None``, the deleting value of the
+        RRset.  The default is ``None``.
 
         *create*, a ``bool``.  If ``True``, create the RRset if it is not found.
         The created RRset is appended to *section*.
@@ -378,6 +380,10 @@ class Message:
         already.  The default is ``False``.  This is useful when creating
         DDNS Update messages, as order matters for them.
 
+        *idna_codec*, a ``dns.name.IDNACodec``, specifies the IDNA
+        encoder/decoder.  If ``None``, the default IDNA 2003 encoder/decoder
+        is used.
+
         Raises ``KeyError`` if the RRset was not found and create was
         ``False``.
 
@@ -386,10 +392,20 @@ class Message:
 
         if isinstance(section, int):
             section_number = section
-            the_section = self.section_from_number(section_number)
+            section = self.section_from_number(section_number)
+        elif isinstance(section, str):
+            section_number = MessageSection.from_text(section)
+            section = self.section_from_number(section_number)
         else:
             section_number = self.section_number(section)
-            the_section = section
+            section = section
+        if isinstance(name, str):
+            name = dns.name.from_text(name, idna_codec=idna_codec)
+        rdtype = dns.rdatatype.RdataType.make(rdtype)
+        rdclass = dns.rdataclass.RdataClass.make(rdclass)
+        covers = dns.rdatatype.RdataType.make(covers)
+        if deleting is not None:
+            deleting = dns.rdataclass.RdataClass.make(deleting)
         key = (section_number, name, rdclass, rdtype, covers, deleting)
         if not force_unique:
             if self.index is not None:
@@ -397,13 +413,13 @@ class Message:
                 if rrset is not None:
                     return rrset
             else:
-                for rrset in the_section:
+                for rrset in section:
                     if rrset.full_match(name, rdclass, rdtype, covers, deleting):
                         return rrset
         if not create:
             raise KeyError
         rrset = dns.rrset.RRset(name, rdclass, rdtype, covers, deleting)
-        the_section.append(rrset)
+        section.append(rrset)
         if self.index is not None:
             self.index[key] = rrset
         return rrset
@@ -418,29 +434,31 @@ class Message:
         deleting: Optional[dns.rdataclass.RdataClass] = None,
         create: bool = False,
         force_unique: bool = False,
+        idna_codec: Optional[dns.name.IDNACodec] = None,
     ) -> Optional[dns.rrset.RRset]:
         """Get the RRset with the given attributes in the specified section.
 
         If the RRset is not found, None is returned.
 
-        *section*, an ``int`` section number, or one of the section
-        attributes of this message.  This specifies the
+        *section*, an ``int`` section number, a ``str`` section name, or one of
+        the section attributes of this message.  This specifies the
         the section of the message to search.  For example::
 
             my_message.get_rrset(my_message.answer, name, rdclass, rdtype)
             my_message.get_rrset(dns.message.ANSWER, name, rdclass, rdtype)
+            my_message.get_rrset("ANSWER", name, rdclass, rdtype)
 
-        *name*, a ``dns.name.Name``, the name of the RRset.
+        *name*, a ``dns.name.Name`` or ``str``, the name of the RRset.
 
-        *rdclass*, an ``int``, the class of the RRset.
+        *rdclass*, an ``int`` or ``str``, the class of the RRset.
 
-        *rdtype*, an ``int``, the type of the RRset.
+        *rdtype*, an ``int`` or ``str``, the type of the RRset.
 
-        *covers*, an ``int`` or ``None``, the covers value of the RRset.
-        The default is ``None``.
+        *covers*, an ``int`` or ``str``, the covers value of the RRset.
+        The default is ``dns.rdatatype.NONE``.
 
-        *deleting*, an ``int`` or ``None``, the deleting value of the RRset.
-        The default is ``None``.
+        *deleting*, an ``int``, ``str``, or ``None``, the deleting value of the
+        RRset.  The default is ``None``.
 
         *create*, a ``bool``.  If ``True``, create the RRset if it is not found.
         The created RRset is appended to *section*.
@@ -450,6 +468,10 @@ class Message:
         already.  The default is ``False``.  This is useful when creating
         DDNS Update messages, as order matters for them.
 
+        *idna_codec*, a ``dns.name.IDNACodec``, specifies the IDNA
+        encoder/decoder.  If ``None``, the default IDNA 2003 encoder/decoder
+        is used.
+
         Returns a ``dns.rrset.RRset object`` or ``None``.
         """
 

diff --git a/tests/test_message.py b/tests/test_message.py
index 29632833c2bc0d600234cae4edb8ce2e9739f5c4..0eb34d7cf33c49c21bcf30eb70e43827fa06a6c7 100644 (file)
--- a/tests/test_message.py
+++ b/tests/test_message.py
@@ -285,6 +285,14 @@ class MessageTestCase(unittest.TestCase):
         self.assertTrue(rrs1 is None)
         self.assertEqual(rrs1, rrs2)
 
+    def test_GetRRsetStrings(self):
+        a = dns.message.from_text(answer_text)
+        a.index = None
+        n = dns.name.from_text("dnspython.org.")
+        rrs1 = a.get_rrset(a.answer, n, dns.rdataclass.IN, dns.rdatatype.SOA)
+        rrs2 = a.get_rrset("ANSWER", "dnspython.org.", "IN", "SOA")
+        self.assertEqual(rrs1, rrs2)
+
     def test_CleanTruncated(self):
         def bad():
             a = dns.message.from_text(answer_text)