Merge pull request #473 from bwelling/enum

WIP: Convert constants to enums

dns/dnssec.py

@@ -41,66 +41,26 @@ class ValidationFailure(dns.exception.DNSException):
     """The DNSSEC signature is invalid."""
 
 
-#: RSAMD5
-RSAMD5 = 1
-#: DH
-DH = 2
-#: DSA
-DSA = 3
-#: ECC
-ECC = 4
-#: RSASHA1
-RSASHA1 = 5
-#: DSANSEC3SHA1
-DSANSEC3SHA1 = 6
-#: RSASHA1NSEC3SHA1
-RSASHA1NSEC3SHA1 = 7
-#: RSASHA256
-RSASHA256 = 8
-#: RSASHA512
-RSASHA512 = 10
-#: ECC-GOST
-ECCGOST = 12
-#: ECDSAP256SHA256
-ECDSAP256SHA256 = 13
-#: ECDSAP384SHA384
-ECDSAP384SHA384 = 14
-#: ED25519
-ED25519 = 15
-#: ED448
-ED448 = 16
-#: INDIRECT
-INDIRECT = 252
-#: PRIVATEDNS
-PRIVATEDNS = 253
-#: PRIVATEOID
-PRIVATEOID = 254
-
-_algorithm_by_text = {
-    'RSAMD5': RSAMD5,
-    'DH': DH,
-    'DSA': DSA,
-    'ECC': ECC,
-    'RSASHA1': RSASHA1,
-    'DSANSEC3SHA1': DSANSEC3SHA1,
-    'RSASHA1NSEC3SHA1': RSASHA1NSEC3SHA1,
-    'RSASHA256': RSASHA256,
-    'RSASHA512': RSASHA512,
-    'ECCGOST': ECCGOST,
-    'ECDSAP256SHA256': ECDSAP256SHA256,
-    'ECDSAP384SHA384': ECDSAP384SHA384,
-    'ED25519': ED25519,
-    'ED448': ED448,
-    'INDIRECT': INDIRECT,
-    'PRIVATEDNS': PRIVATEDNS,
-    'PRIVATEOID': PRIVATEOID,
-}
-
-# 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.
-
-_algorithm_by_value = {y: x for x, y in _algorithm_by_text.items()}
+class Algorithm(enum.IntEnum):
+    RSAMD5 = 1
+    DH = 2
+    DSA = 3
+    ECC = 4
+    RSASHA1 = 5
+    DSANSEC3SHA1 = 6
+    RSASHA1NSEC3SHA1 = 7
+    RSASHA256 = 8
+    RSASHA512 = 10
+    ECCGOST = 12
+    ECDSAP256SHA256 = 13
+    ECDSAP384SHA384 = 14
+    ED25519 = 15
+    ED448 = 16
+    INDIRECT = 252
+    PRIVATEDNS = 253
+    PRIVATEOID = 254
+
+globals().update(Algorithm.__members__)
 
 
 def algorithm_from_text(text):
@@ -111,10 +71,10 @@ def algorithm_from_text(text):
     Returns an ``int``.
     """
 
-    value = _algorithm_by_text.get(text.upper())
-    if value is None:
-        value = int(text)
-    return value
+    try:
+        return Algorithm[text.upper()]
+    except KeyError:
+        return int(text)
 
 
 def algorithm_to_text(value):
@@ -125,10 +85,10 @@ def algorithm_to_text(value):
     Returns a ``str``, the name of a DNSSEC algorithm.
     """
 
-    text = _algorithm_by_value.get(value)
-    if text is None:
-        text = str(value)
-    return text
+    try:
+        return Algorithm(value).name
+    except ValueError:
+        return str(value)
 
 
 def _to_rdata(record, origin):

dns/flags.py

@@ -17,76 +17,52 @@
 
 """DNS Message Flags."""
 
-# Standard DNS flags
-
-#: Query Response
-QR = 0x8000
-#: Authoritative Answer
-AA = 0x0400
-#: Truncated Response
-TC = 0x0200
-#: Recursion Desired
-RD = 0x0100
-#: Recursion Available
-RA = 0x0080
-#: Authentic Data
-AD = 0x0020
-#: Checking Disabled
-CD = 0x0010
-
-# EDNS flags
-
-#: DNSSEC answer OK
-DO = 0x8000
+import enum
 
-_by_text = {
-    'QR': QR,
-    'AA': AA,
-    'TC': TC,
-    'RD': RD,
-    'RA': RA,
-    'AD': AD,
-    'CD': CD
-}
-
-_edns_by_text = {
-    'DO': DO
-}
-
-
-# We construct the inverse mappings programmatically to ensure that we
-# cannot make any mistakes (e.g. omissions, cut-and-paste errors) that
-# would cause the mappings not to be true inverses.
+# Standard DNS flags
 
-_by_value = {y: x for x, y in _by_text.items()}
+class Flag(enum.IntFlag):
+    #: Query Response
+    QR = 0x8000
+    #: Authoritative Answer
+    AA = 0x0400
+    #: Truncated Response
+    TC = 0x0200
+    #: Recursion Desired
+    RD = 0x0100
+    #: Recursion Available
+    RA = 0x0080
+    #: Authentic Data
+    AD = 0x0020
+    #: Checking Disabled
+    CD = 0x0010
+
+globals().update(Flag.__members__)
 
-_edns_by_value = {y: x for x, y in _edns_by_text.items()}
 
+# EDNS flags
 
-def _order_flags(table):
-    order = list(table.items())
-    order.sort()
-    order.reverse()
-    return order
+class EDNSFlag(enum.IntFlag):
+    #: DNSSEC answer OK
+    DO = 0x8000
 
-_flags_order = _order_flags(_by_value)
 
-_edns_flags_order = _order_flags(_edns_by_value)
+globals().update(EDNSFlag.__members__)
 
 
-def _from_text(text, table):
+def _from_text(text, enum_class):
     flags = 0
     tokens = text.split()
     for t in tokens:
-        flags = flags | table[t.upper()]
+        flags |= enum_class[t.upper()]
     return flags
 
 
-def _to_text(flags, table, order):
+def _to_text(flags, enum_class):
     text_flags = []
-    for k, v in order:
-        if flags & k != 0:
-            text_flags.append(v)
+    for k, v in enum_class.__members__.items():
+        if flags & v != 0:
+            text_flags.append(k)
     return ' '.join(text_flags)
 
 
@@ -97,7 +73,7 @@ def from_text(text):
     Returns an ``int``
     """
 
-    return _from_text(text, _by_text)
+    return _from_text(text, Flag)
 
 
 def to_text(flags):
@@ -107,7 +83,7 @@ def to_text(flags):
     Returns a ``str``.
     """
 
-    return _to_text(flags, _by_value, _flags_order)
+    return _to_text(flags, Flag)
 
 
 def edns_from_text(text):
@@ -117,7 +93,7 @@ def edns_from_text(text):
     Returns an ``int``
     """
 
-    return _from_text(text, _edns_by_text)
+    return _from_text(text, EDNSFlag)
 
 
 def edns_to_text(flags):
@@ -127,4 +103,4 @@ def edns_to_text(flags):
     Returns a ``str``.
     """
 
-    return _to_text(flags, _edns_by_value, _edns_flags_order)
+    return _to_text(flags, EDNSFlag)

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,

dns/opcode.py

@@ -17,32 +17,23 @@
 
 """DNS Opcodes."""
 
-import dns.exception
-
-#: Query
-QUERY = 0
-#: Inverse Query (historical)
-IQUERY = 1
-#: Server Status (unspecified and unimplemented anywhere)
-STATUS = 2
-#: Notify
-NOTIFY = 4
-#: Dynamic Update
-UPDATE = 5
+import enum
 
-_by_text = {
-    'QUERY': QUERY,
-    'IQUERY': IQUERY,
-    'STATUS': STATUS,
-    'NOTIFY': NOTIFY,
-    'UPDATE': UPDATE
-}
+import dns.exception
 
-# 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.
+class Opcode(enum.IntEnum):
+    #: Query
+    QUERY = 0
+    #: Inverse Query (historical)
+    IQUERY = 1
+    #: Server Status (unspecified and unimplemented anywhere)
+    STATUS = 2
+    #: Notify
+    NOTIFY = 4
+    #: Dynamic Update
+    UPDATE = 5
 
-_by_value = {y: x for x, y in _by_text.items()}
+globals().update(Opcode.__members__)
 
 
 class UnknownOpcode(dns.exception.DNSException):
@@ -62,11 +53,14 @@ def from_text(text):
     if text.isdigit():
         value = int(text)
         if value >= 0 and value <= 15:
-            return value
-    value = _by_text.get(text.upper())
-    if value is None:
+            try:
+                return Opcode(value)
+            except ValueError:
+                return value
+    try:
+        return Opcode[text.upper()]
+    except:
         raise UnknownOpcode
-    return value
 
 
 def from_flags(flags):
@@ -102,10 +96,10 @@ def to_text(value):
     Returns a ``str``.
     """
 
-    text = _by_value.get(value)
-    if text is None:
-        text = str(value)
-    return text
+    try:
+        return Opcode(value).name
+    except ValueError:
+        return str(value)
 
 
 def is_update(flags):

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',