section enums

diff --git a/dns/message.py b/dns/message.py
index 132149dba45d1423a597eb5a7242104c8840aed1..638f2048c954ecfa03315588e1313856de99a382 100644 (file)
--- a/dns/message.py
+++ b/dns/message.py
@@ -23,6 +23,7 @@ import struct
 import time
 
 import dns.edns
+import dns.enum
 import dns.exception
 import dns.flags
 import dns.name
@@ -77,22 +78,20 @@ class Truncated(dns.exception.DNSException):
         """
         return self.kwargs['message']
 
+class MessageSection(dns.enum.IntEnum):
+    """Message sections"""
+    QUESTION = 0
+    ANSWER = 1
+    AUTHORITY = 2
+    ADDITIONAL = 3
 
-#: The question section number
-QUESTION = 0
-
-#: The answer section number
-ANSWER = 1
-
-#: The authority section number
-AUTHORITY = 2
-
-#: The additional section number
-ADDITIONAL = 3
+globals().update(MessageSection.__members__)
 
 class Message:
     """A DNS message."""
 
+    _section_enum = MessageSection
+
     def __init__(self, id=None):
         if id is None:
             self.id = dns.entropy.random_16()
@@ -155,32 +154,14 @@ class Message:
             s.write('payload %d\n' % self.payload)
         for opt in self.options:
             s.write('option %s\n' % opt.to_text())
-        is_update = dns.opcode.is_update(self.flags)
-        if is_update:
-            s.write(';ZONE\n')
-        else:
-            s.write(';QUESTION\n')
-        for rrset in self.question:
-            s.write(rrset.to_text(origin, relativize, **kw))
-            s.write('\n')
-        if is_update:
-            s.write(';PREREQ\n')
-        else:
-            s.write(';ANSWER\n')
-        for rrset in self.answer:
-            s.write(rrset.to_text(origin, relativize, **kw))
-            s.write('\n')
-        if is_update:
-            s.write(';UPDATE\n')
-        else:
-            s.write(';AUTHORITY\n')
-        for rrset in self.authority:
-            s.write(rrset.to_text(origin, relativize, **kw))
-            s.write('\n')
-        s.write(';ADDITIONAL\n')
-        for rrset in self.additional:
+
             s.write(rrset.to_text(origin, relativize, **kw))
             s.write('\n')
+        for (name, which) in self._section_enum.__members__.items():
+            s.write(f';{name}\n')
+            for rrset in self.section_from_number(which):
+                s.write(rrset.to_text(origin, relativize, **kw))
+                s.write('\n')
         #
         # We strip off the final \n so the caller can print the result without
         # doing weird things to get around eccentricities in Python print

diff --git a/dns/update.py b/dns/update.py
index 951d895109904c3eaaab9c6f142c7f118b35e1f0..4878db715b8c1d0dc2db0320340cc65d6971a32d 100644 (file)
--- a/dns/update.py
+++ b/dns/update.py
@@ -27,8 +27,18 @@ import dns.rdataset
 import dns.tsig
 
 
+class UpdateSection(dns.enum.IntEnum):
+    """Update sections"""
+    ZONE = 0
+    PREREQ = 1
+    UPDATE = 2
+    ADDITIONAL = 3
+
+
 class Update(dns.message.Message):
 
+    _section_enum = UpdateSection
+
     def __init__(self, zone, rdclass=dns.rdataclass.IN, keyring=None,
                  keyname=None, keyalgorithm=dns.tsig.default_algorithm):
         """Initialize a new DNS Update object.