From f054938849c4c8d651fc0742e6889922c74b7d85 Mon Sep 17 00:00:00 2001
From: Bob Halley <halley@dnspython.org>
Date: Tue, 17 Jul 2018 07:02:38 -0700
Subject: [PATCH] Add a way to dynamically register an rdata module.

---
 dns/rdata.py         | 32 ++++++++++++++++++++++++++++++++
 dns/rdatatype.py     | 17 +++++++++++++++++
 tests/test_rdata.py  | 16 ++++++++++++++++
 tests/ttxt_module.py |  4 ++++
 4 files changed, 69 insertions(+)
 create mode 100644 tests/ttxt_module.py

diff --git a/dns/rdata.py b/dns/rdata.py
index 89b78492..6f8a3983 100644
--- a/dns/rdata.py
+++ b/dns/rdata.py
@@ -417,3 +417,35 @@ def from_wire(rdclass, rdtype, wire, current, rdlen, origin=None):
     wire = dns.wiredata.maybe_wrap(wire)
     cls = get_rdata_class(rdclass, rdtype)
     return cls.from_wire(rdclass, rdtype, wire, current, rdlen, origin)
+
+
+class RdatatypeExists(dns.exception.DNSException):
+    """DNS rdatatype already exists."""
+    supp_kwargs = set(['rdclass', 'rdtype'])
+    fmt = "The rdata type with class {rdclass} and rdtype {rdtype} " + \
+        "already exists."
+
+
+def register_type(implementation, rdtype, rdtype_text, is_singleton=False,
+                  rdclass=dns.rdataclass.IN):
+    """Dynamically register a module to handle an rdatatype.
+
+    *implementation*, a module implementing the type in the usual dnspython
+    way.
+
+    *rdtype*, an ``int``, the rdatatype to register.
+
+    *rdtype_text*, a ``text``, the textual form of the rdatatype.
+
+    *is_singleton*, a ``bool``, indicating if the type is a singleton (i.e.
+    RRsets of the type can have only one member.)
+
+    *rdclass*, the rdataclass of the type, or ``dns.rdataclass.ANY`` if
+    it applies to all classes.
+    """
+
+    existing_cls = get_rdata_class(rdclass, rdtype)
+    if existing_cls != GenericRdata:
+        raise RdatatypeExists(rdclass=rdclass, rdtype=rdtype)
+    _rdata_modules[(rdclass, rdtype)] = implementation
+    dns.rdatatype.register_type(rdtype, rdtype_text, is_singleton)
diff --git a/dns/rdatatype.py b/dns/rdatatype.py
index 48d2aa67..6d253588 100644
--- a/dns/rdatatype.py
+++ b/dns/rdatatype.py
@@ -266,3 +266,20 @@ def is_singleton(rdtype):
     if rdtype in _singletons:
         return True
     return False
+
+
+def register_type(rdtype, rdtype_text, is_singleton=False):
+    """Dynamically register an rdatatype.
+
+    *rdtype*, an ``int``, the rdatatype to register.
+
+    *rdtype_text*, a ``text``, the textual form of the rdatatype.
+
+    *is_singleton*, a ``bool``, indicating if the type is a singleton (i.e.
+    RRsets of the type can have only one member.)
+    """
+
+    _by_text[rdtype_text] = rdtype
+    _by_value[rdtype] = rdtype_text
+    if is_singleton:
+        _singletons[rdtype] = True
diff --git a/tests/test_rdata.py b/tests/test_rdata.py
index fbf11b39..b67470e1 100644
--- a/tests/test_rdata.py
+++ b/tests/test_rdata.py
@@ -22,6 +22,8 @@ import dns.rdata
 import dns.rdataclass
 import dns.rdatatype
 
+import ttxt_module
+
 class RdataTestCase(unittest.TestCase):
 
     def test_str(self):
@@ -32,5 +34,19 @@ class RdataTestCase(unittest.TestCase):
         rdata = dns.rdata.from_text(dns.rdataclass.IN, dns.rdatatype.A, u"1.2.3.4")
         self.failUnless(rdata.address == "1.2.3.4")
 
+    def test_module_registration(self):
+        TTXT = 64001
+        dns.rdata.register_type(ttxt_module, TTXT, 'TTXT')
+        rdata = dns.rdata.from_text(dns.rdataclass.IN, TTXT, 'hello world')
+        self.failUnless(rdata.strings == [b'hello', b'world'])
+        self.failUnless(dns.rdatatype.to_text(TTXT) == 'TTXT')
+        self.failUnless(dns.rdatatype.from_text('TTXT') == TTXT)
+
+    def test_module_reregistration(self):
+        def bad():
+            TTXTTWO = dns.rdatatype.TXT
+            dns.rdata.register_type(ttxt_module, TTXTTWO, 'TTXTTWO')
+        self.failUnlessRaises(dns.rdata.RdatatypeExists, bad)
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/ttxt_module.py b/tests/ttxt_module.py
new file mode 100644
index 00000000..c66131bb
--- /dev/null
+++ b/tests/ttxt_module.py
@@ -0,0 +1,4 @@
+import dns.rdtypes.txtbase
+
+class TTXT(dns.rdtypes.txtbase.TXTBase):
+    """Test TXT-like record"""
-- 
2.47.3