From: Brian Wellington <bwelling@xbill.org>
Date: Mon, 18 May 2020 17:25:11 +0000 (-0700)
Subject: Start converting rdatatype/rdataclass to enum.
X-Git-Tag: v2.0.0rc1~198^2~5
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=65e107bd0d43a90b0d0d016b594f6587e2875132;p=thirdparty%2Fdnspython.git

Start converting rdatatype/rdataclass to enum.
---

diff --git a/dns/message.py b/dns/message.py
index 44802807..3446d1fc 100644
--- a/dns/message.py
+++ b/dns/message.py
@@ -1110,10 +1110,8 @@ def make_query(qname, rdtype, rdclass=dns.rdataclass.IN, use_edns=None,
 
     if isinstance(qname, str):
         qname = dns.name.from_text(qname, idna_codec=idna_codec)
-    if isinstance(rdtype, str):
-        rdtype = dns.rdatatype.from_text(rdtype)
-    if isinstance(rdclass, str):
-        rdclass = dns.rdataclass.from_text(rdclass)
+    rdtype = dns.rdatatype.to_enum(rdtype)
+    rdclass = dns.rdataclass.to_enum(rdclass)
     m = Message()
     m.flags |= dns.flags.RD
     m.find_rrset(m.question, qname, rdclass, rdtype, create=True,
diff --git a/dns/query.py b/dns/query.py
index 34adb62d..de5c300b 100644
--- a/dns/query.py
+++ b/dns/query.py
@@ -817,8 +817,7 @@ def xfr(where, zone, rdtype=dns.rdatatype.AXFR, rdclass=dns.rdataclass.IN,
 
     if isinstance(zone, str):
         zone = dns.name.from_text(zone)
-    if isinstance(rdtype, str):
-        rdtype = dns.rdatatype.from_text(rdtype)
+    rdtype = dns.rdatatype.to_enum(rdtype)
     q = dns.message.make_query(zone, rdtype, rdclass)
     if rdtype == dns.rdatatype.IXFR:
         rrset = dns.rrset.from_text(zone, 0, 'IN', 'SOA',
diff --git a/dns/rdata.py b/dns/rdata.py
index 69f08ec2..fd42e4bc 100644
--- a/dns/rdata.py
+++ b/dns/rdata.py
@@ -433,6 +433,8 @@ def from_text(rdclass, rdtype, tok, origin=None, relativize=True,
 
     if isinstance(tok, str):
         tok = dns.tokenizer.Tokenizer(tok, idna_codec=idna_codec)
+    rdclass = dns.rdataclass.to_enum(rdclass)
+    rdtype = dns.rdatatype.to_enum(rdtype)
     cls = get_rdata_class(rdclass, rdtype)
     if cls != GenericRdata:
         # peek at first token
@@ -482,6 +484,8 @@ def from_wire(rdclass, rdtype, wire, current, rdlen, origin=None):
     """
 
     wire = dns.wiredata.maybe_wrap(wire)
+    rdclass = dns.rdataclass.to_enum(rdclass)
+    rdtype = dns.rdatatype.to_enum(rdtype)
     cls = get_rdata_class(rdclass, rdtype)
     return cls.from_wire(rdclass, rdtype, wire, current, rdlen, origin)
 
diff --git a/dns/rdataclass.py b/dns/rdataclass.py
index b88aa85b..bc883fd2 100644
--- a/dns/rdataclass.py
+++ b/dns/rdataclass.py
@@ -17,40 +17,25 @@
 
 """DNS Rdata Classes."""
 
+import enum
 import re
 
 import dns.exception
 
-RESERVED0 = 0
-IN = 1
-CH = 3
-HS = 4
-NONE = 254
-ANY = 255
-
-_by_text = {
-    'RESERVED0': RESERVED0,
-    'IN': IN,
-    'CH': CH,
-    'HS': HS,
-    'NONE': NONE,
-    'ANY': ANY
-}
-
-# We construct the inverse mapping programmatically to ensure that we
-# cannot make any mistakes (e.g. omissions, cut-and-paste errors) that
-# would cause the mapping not to be true inverse.
-
-_by_value = {y: x for x, y in _by_text.items()}
-
-# Now that we've built the inverse map, we can add class aliases to
-# the _by_text mapping.
-
-_by_text.update({
-    'INTERNET': IN,
-    'CHAOS': CH,
-    'HESIOD': HS
-})
+class RdataClass(enum.IntEnum):
+    """DNS Rdata Class"""
+    RESERVED0 = 0
+    IN = 1
+    INTERNET = IN
+    CH = 3
+    CHAOS = CH
+    HS = 4
+    HESIOD = HS
+    NONE = 254
+    ANY = 255
+
+for (name, value) in RdataClass.__members__.items():
+    globals()[name] = value
 
 _metaclasses = {
     NONE: True,
@@ -79,8 +64,9 @@ def from_text(text):
     Returns an ``int``.
     """
 
-    value = _by_text.get(text.upper())
-    if value is None:
+    try:
+        value = RdataClass[text.upper()]
+    except KeyError:
         match = _unknown_class_pattern.match(text)
         if match is None:
             raise UnknownRdataclass
@@ -91,7 +77,7 @@ def from_text(text):
 
 
 def to_text(value):
-    """Convert a DNS rdata type value to text.
+    """Convert a DNS rdata class value to text.
 
     If the value has a known mnemonic, it will be used, otherwise the
     DNS generic class syntax will be used.
@@ -103,10 +89,28 @@ def to_text(value):
 
     if value < 0 or value > 65535:
         raise ValueError("class must be between >= 0 and <= 65535")
-    text = _by_value.get(value)
-    if text is None:
-        text = 'CLASS' + repr(value)
-    return text
+    try:
+        return RdataClass(value).name
+    except ValueError:
+        return f'CLASS{value}'
+
+
+def to_enum(value):
+    """Convert a DNS rdata class value to an enumerated type, if possible.
+
+    *value*, an ``int`` or ``str``, the rdata class.
+
+    Returns an ``int``.
+    """
+
+    if isinstance(value, str):
+        return from_text(value)
+    if value < 0 or value > 65535:
+        raise ValueError("class must be between >= 0 and <= 65535")
+    try:
+        return RdataClass(value)
+    except ValueError:
+        return value
 
 
 def is_metaclass(rdclass):
diff --git a/dns/rdataset.py b/dns/rdataset.py
index 9d6a6a8e..1661ca6f 100644
--- a/dns/rdataset.py
+++ b/dns/rdataset.py
@@ -309,10 +309,8 @@ def from_text_list(rdclass, rdtype, ttl, text_rdatas, idna_codec=None):
     Returns a ``dns.rdataset.Rdataset`` object.
     """
 
-    if isinstance(rdclass, str):
-        rdclass = dns.rdataclass.from_text(rdclass)
-    if isinstance(rdtype, str):
-        rdtype = dns.rdatatype.from_text(rdtype)
+    rdclass = dns.rdataclass.to_enum(rdclass)
+    rdtype = dns.rdatatype.to_enum(rdtype)
     r = Rdataset(rdclass, rdtype)
     r.update_ttl(ttl)
     for t in text_rdatas:
diff --git a/dns/rdatatype.py b/dns/rdatatype.py
index 150a047f..24442be7 100644
--- a/dns/rdatatype.py
+++ b/dns/rdatatype.py
@@ -17,164 +17,95 @@
 
 """DNS Rdata Types."""
 
+import enum
 import re
 
 import dns.exception
 
-NONE = 0
-A = 1
-NS = 2
-MD = 3
-MF = 4
-CNAME = 5
-SOA = 6
-MB = 7
-MG = 8
-MR = 9
-NULL = 10
-WKS = 11
-PTR = 12
-HINFO = 13
-MINFO = 14
-MX = 15
-TXT = 16
-RP = 17
-AFSDB = 18
-X25 = 19
-ISDN = 20
-RT = 21
-NSAP = 22
-NSAP_PTR = 23
-SIG = 24
-KEY = 25
-PX = 26
-GPOS = 27
-AAAA = 28
-LOC = 29
-NXT = 30
-SRV = 33
-NAPTR = 35
-KX = 36
-CERT = 37
-A6 = 38
-DNAME = 39
-OPT = 41
-APL = 42
-DS = 43
-SSHFP = 44
-IPSECKEY = 45
-RRSIG = 46
-NSEC = 47
-DNSKEY = 48
-DHCID = 49
-NSEC3 = 50
-NSEC3PARAM = 51
-TLSA = 52
-HIP = 55
-NINFO = 56
-CDS = 59
-CDNSKEY = 60
-OPENPGPKEY = 61
-CSYNC = 62
-SPF = 99
-UNSPEC = 103
-EUI48 = 108
-EUI64 = 109
-TKEY = 249
-TSIG = 250
-IXFR = 251
-AXFR = 252
-MAILB = 253
-MAILA = 254
-ANY = 255
-URI = 256
-CAA = 257
-AVC = 258
-TA = 32768
-DLV = 32769
-
-_by_text = {
-    'NONE': NONE,
-    'A': A,
-    'NS': NS,
-    'MD': MD,
-    'MF': MF,
-    'CNAME': CNAME,
-    'SOA': SOA,
-    'MB': MB,
-    'MG': MG,
-    'MR': MR,
-    'NULL': NULL,
-    'WKS': WKS,
-    'PTR': PTR,
-    'HINFO': HINFO,
-    'MINFO': MINFO,
-    'MX': MX,
-    'TXT': TXT,
-    'RP': RP,
-    'AFSDB': AFSDB,
-    'X25': X25,
-    'ISDN': ISDN,
-    'RT': RT,
-    'NSAP': NSAP,
-    'NSAP-PTR': NSAP_PTR,
-    'SIG': SIG,
-    'KEY': KEY,
-    'PX': PX,
-    'GPOS': GPOS,
-    'AAAA': AAAA,
-    'LOC': LOC,
-    'NXT': NXT,
-    'SRV': SRV,
-    'NAPTR': NAPTR,
-    'KX': KX,
-    'CERT': CERT,
-    'A6': A6,
-    'DNAME': DNAME,
-    'OPT': OPT,
-    'APL': APL,
-    'DS': DS,
-    'SSHFP': SSHFP,
-    'IPSECKEY': IPSECKEY,
-    'RRSIG': RRSIG,
-    'NSEC': NSEC,
-    'DNSKEY': DNSKEY,
-    'DHCID': DHCID,
-    'NSEC3': NSEC3,
-    'NSEC3PARAM': NSEC3PARAM,
-    'TLSA': TLSA,
-    'HIP': HIP,
-    'NINFO': NINFO,
-    'CDS': CDS,
-    'CDNSKEY': CDNSKEY,
-    'OPENPGPKEY': OPENPGPKEY,
-    'CSYNC': CSYNC,
-    'SPF': SPF,
-    'UNSPEC': UNSPEC,
-    'EUI48': EUI48,
-    'EUI64': EUI64,
-    'TKEY': TKEY,
-    'TSIG': TSIG,
-    'IXFR': IXFR,
-    'AXFR': AXFR,
-    'MAILB': MAILB,
-    'MAILA': MAILA,
-    'ANY': ANY,
-    'URI': URI,
-    'CAA': CAA,
-    'AVC': AVC,
-    'TA': TA,
-    'DLV': DLV,
-}
-
-# We construct the inverse mapping programmatically to ensure that we
-# cannot make any mistakes (e.g. omissions, cut-and-paste errors) that
-# would cause the mapping not to be true inverse.
-
-_by_value = {y: x for x, y in _by_text.items()}
-# Render type 0 as "TYPE0" not "NONE", as NONE is a dnspython-ism and not
-# an official mnemonic.
-_by_value[0] = 'TYPE0'
+class RdataType(enum.IntEnum):
+    """DNS Rdata Type"""
+    TYPE0 = 0
+    NONE = 0
+    A = 1
+    NS = 2
+    MD = 3
+    MF = 4
+    CNAME = 5
+    SOA = 6
+    MB = 7
+    MG = 8
+    MR = 9
+    NULL = 10
+    WKS = 11
+    PTR = 12
+    HINFO = 13
+    MINFO = 14
+    MX = 15
+    TXT = 16
+    RP = 17
+    AFSDB = 18
+    X25 = 19
+    ISDN = 20
+    RT = 21
+    NSAP = 22
+    NSAP_PTR = 23
+    SIG = 24
+    KEY = 25
+    PX = 26
+    GPOS = 27
+    AAAA = 28
+    LOC = 29
+    NXT = 30
+    SRV = 33
+    NAPTR = 35
+    KX = 36
+    CERT = 37
+    A6 = 38
+    DNAME = 39
+    OPT = 41
+    APL = 42
+    DS = 43
+    SSHFP = 44
+    IPSECKEY = 45
+    RRSIG = 46
+    NSEC = 47
+    DNSKEY = 48
+    DHCID = 49
+    NSEC3 = 50
+    NSEC3PARAM = 51
+    TLSA = 52
+    HIP = 55
+    NINFO = 56
+    CDS = 59
+    CDNSKEY = 60
+    OPENPGPKEY = 61
+    CSYNC = 62
+    SPF = 99
+    UNSPEC = 103
+    EUI48 = 108
+    EUI64 = 109
+    TKEY = 249
+    TSIG = 250
+    IXFR = 251
+    AXFR = 252
+    MAILB = 253
+    MAILA = 254
+    ANY = 255
+    URI = 256
+    CAA = 257
+    AVC = 258
+    TA = 32768
+    DLV = 32769
+
+_by_text = {}
+_by_value = {}
+
+for (name, value) in RdataType.__members__.items():
+    globals()[name] = value
+    real_name = name.replace('_', '-')
+    _by_text[real_name] = value
+    if value not in _by_value:
+        _by_value[value] = real_name
 
 _metatypes = {
     OPT: True
@@ -240,6 +171,24 @@ def to_text(value):
     return text
 
 
+def to_enum(value):
+    """Convert a DNS rdata type value to an enumerated type, if possible.
+
+    *value*, an ``int`` or ``str``, the rdata type.
+
+    Returns an ``int``.
+    """
+
+    if isinstance(value, str):
+        return from_text(value)
+    if value < 0 or value > 65535:
+        raise ValueError("type must be between >= 0 and <= 65535")
+    try:
+        return RdataType(value)
+    except ValueError:
+        return value
+
+
 def is_metatype(rdtype):
     """True if the specified type is a metatype.
 
diff --git a/dns/rrset.py b/dns/rrset.py
index d05792a9..864af2af 100644
--- a/dns/rrset.py
+++ b/dns/rrset.py
@@ -140,10 +140,8 @@ def from_text_list(name, ttl, rdclass, rdtype, text_rdatas,
 
     if isinstance(name, str):
         name = dns.name.from_text(name, None, idna_codec=idna_codec)
-    if isinstance(rdclass, str):
-        rdclass = dns.rdataclass.from_text(rdclass)
-    if isinstance(rdtype, str):
-        rdtype = dns.rdatatype.from_text(rdtype)
+    rdclass = dns.rdataclass.to_enum(rdclass)
+    rdtype = dns.rdatatype.to_enum(rdtype)
     r = RRset(name, rdclass, rdtype)
     r.update_ttl(ttl)
     for t in text_rdatas:
diff --git a/dns/update.py b/dns/update.py
index 56471173..74096e1d 100644
--- a/dns/update.py
+++ b/dns/update.py
@@ -58,8 +58,7 @@ class Update(dns.message.Message):
         if isinstance(zone, str):
             zone = dns.name.from_text(zone)
         self.origin = zone
-        if isinstance(rdclass, str):
-            rdclass = dns.rdataclass.from_text(rdclass)
+        rdclass = dns.rdataclass.to_enum(rdclass)
         self.zone_rdclass = rdclass
         self.find_rrset(self.question, self.origin, rdclass, dns.rdatatype.SOA,
                         create=True, force_unique=True)
@@ -109,9 +108,7 @@ class Update(dns.message.Message):
                 for rd in args:
                     self._add_rr(name, ttl, rd, section=section)
             else:
-                rdtype = args.pop(0)
-                if isinstance(rdtype, str):
-                    rdtype = dns.rdatatype.from_text(rdtype)
+                rdtype = dns.rdatatype.to_enum(args.pop(0))
                 if replace:
                     self.delete(name, rdtype)
                 for s in args:
@@ -165,9 +162,7 @@ class Update(dns.message.Message):
                 for rd in args:
                     self._add_rr(name, 0, rd, dns.rdataclass.NONE)
             else:
-                rdtype = args.pop(0)
-                if isinstance(rdtype, str):
-                    rdtype = dns.rdatatype.from_text(rdtype)
+                rdtype = dns.rdatatype.to_enum(args.pop(0))
                 if len(args) == 0:
                     self.find_rrset(self.authority, name,
                                     self.zone_rdclass, rdtype,
@@ -229,9 +224,7 @@ class Update(dns.message.Message):
                 args.insert(0, 0)
             self._add(False, self.answer, name, *args)
         else:
-            rdtype = args[0]
-            if isinstance(rdtype, str):
-                rdtype = dns.rdatatype.from_text(rdtype)
+            rdtype = dns.rdatatype.to_enum(args[0])
             self.find_rrset(self.answer, name,
                             dns.rdataclass.ANY, rdtype,
                             dns.rdatatype.NONE, None,
@@ -249,8 +242,7 @@ class Update(dns.message.Message):
                             dns.rdatatype.NONE, None,
                             True, True)
         else:
-            if isinstance(rdtype, str):
-                rdtype = dns.rdatatype.from_text(rdtype)
+            rdtype = dns.rdatatype.to_enum(rdtype)
             self.find_rrset(self.answer, name,
                             dns.rdataclass.NONE, rdtype,
                             dns.rdatatype.NONE, None,
diff --git a/dns/zone.py b/dns/zone.py
index be7333c3..a75b6b26 100644
--- a/dns/zone.py
+++ b/dns/zone.py
@@ -270,10 +270,9 @@ class Zone(object):
         """
 
         name = self._validate_name(name)
-        if isinstance(rdtype, str):
-            rdtype = dns.rdatatype.from_text(rdtype)
-        if isinstance(covers, str):
-            covers = dns.rdatatype.from_text(covers)
+        rdtype = dns.rdatatype.to_enum(rdtype)
+        if covers is not None:
+            covers = dns.rdatatype.to_enum(covers)
         node = self.find_node(name, create)
         return node.find_rdataset(self.rdclass, rdtype, covers, create)
 
@@ -349,10 +348,9 @@ class Zone(object):
         """
 
         name = self._validate_name(name)
-        if isinstance(rdtype, str):
-            rdtype = dns.rdatatype.from_text(rdtype)
-        if isinstance(covers, str):
-            covers = dns.rdatatype.from_text(covers)
+        rdtype = dns.rdatatype.to_enum(rdtype)
+        if covers is not None:
+            covers = dns.rdatatype.to_enum(covers)
         node = self.get_node(name)
         if node is not None:
             node.delete_rdataset(self.rdclass, rdtype, covers)
@@ -423,10 +421,9 @@ class Zone(object):
         """
 
         name = self._validate_name(name)
-        if isinstance(rdtype, str):
-            rdtype = dns.rdatatype.from_text(rdtype)
-        if isinstance(covers, str):
-            covers = dns.rdatatype.from_text(covers)
+        rdtype = dns.rdatatype.to_enum(rdtype)
+        if covers is not None:
+            covers = dns.rdatatype.to_enum(covers)
         rdataset = self.nodes[name].find_rdataset(self.rdclass, rdtype, covers)
         rrset = dns.rrset.RRset(name, self.rdclass, rdtype, covers)
         rrset.update(rdataset)
@@ -496,10 +493,9 @@ class Zone(object):
         RRSIG rdataset.
         """
 
-        if isinstance(rdtype, str):
-            rdtype = dns.rdatatype.from_text(rdtype)
-        if isinstance(covers, str):
-            covers = dns.rdatatype.from_text(covers)
+        rdtype = dns.rdatatype.to_enum(rdtype)
+        if covers is not None:
+            covers = dns.rdatatype.to_enum(covers)
         for (name, node) in self.items():
             for rds in node:
                 if rdtype == dns.rdatatype.ANY or \
@@ -526,10 +522,9 @@ class Zone(object):
         RRSIG rdataset.
         """
 
-        if isinstance(rdtype, str):
-            rdtype = dns.rdatatype.from_text(rdtype)
-        if isinstance(covers, str):
-            covers = dns.rdatatype.from_text(covers)
+        rdtype = dns.rdatatype.to_enum(rdtype)
+        if covers is not None:
+            covers = dns.rdatatype.to_enum(covers)
         for (name, node) in self.items():
             for rds in node:
                 if rdtype == dns.rdatatype.ANY or \