Convert message flags to enums.

1 files changed, 34 insertions, 58 deletions

diff --git a/dns/flags.py b/dns/flags.py
index e5fdd07..4eb6d90 100644
--- a/dns/flags.py
+++ b/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)