Skip to content

Commit

Permalink
Support registering new types with classes. (#1167)
Browse files Browse the repository at this point in the history
* Support registering new types with classes.

Previously, dns.rdata.register_type() required passing a module which
contained the implementation of the new type, and it would extract the
class from the module.  This change allows passing the class directly.
  • Loading branch information
bwelling authored Nov 29, 2024
1 parent e6373b6 commit f972afd
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 5 deletions.
12 changes: 7 additions & 5 deletions dns/rdata.py
Original file line number Diff line number Diff line change
Expand Up @@ -891,8 +891,8 @@ def register_type(
) -> None:
"""Dynamically register a module to handle an rdatatype.
*implementation*, a module implementing the type in the usual dnspython
way.
*implementation*, a subclass of ``dns.rdata.Rdata`` implementing the type,
or a module containing such a class named by its text form.
*rdtype*, an ``int``, the rdatatype to register.
Expand All @@ -909,7 +909,9 @@ def register_type(
existing_cls = get_rdata_class(rdclass, rdtype)
if existing_cls != GenericRdata or dns.rdatatype.is_metatype(rdtype):
raise RdatatypeExists(rdclass=rdclass, rdtype=rdtype)
_rdata_classes[(rdclass, rdtype)] = getattr(
implementation, rdtype_text.replace("-", "_")
)
if isinstance(implementation, type) and issubclass(implementation, Rdata):
impclass = implementation
else:
impclass = getattr(implementation, rdtype_text.replace("-", "_"))
_rdata_classes[(rdclass, rdtype)] = impclass
dns.rdatatype.register_type(rdtype, rdtype_text, is_singleton)
14 changes: 14 additions & 0 deletions tests/test_rdata.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,20 @@ def test_module_registration(self):
self.assertEqual(dns.rdatatype.from_text("ttxt"), TTXT)
self.assertEqual(dns.rdatatype.RdataType.make("ttxt"), TTXT)

def test_class_registration(self):
CTXT = 64003
class CTXTImp(dns.rdtypes.txtbase.TXTBase):
"""Test TXT-like record"""

dns.rdata.register_type(CTXTImp, CTXT, "CTXT")
rdata = dns.rdata.from_text(dns.rdataclass.IN, CTXT, "hello world")
self.assertEqual(rdata.strings, (b"hello", b"world"))
self.assertEqual(dns.rdatatype.to_text(CTXT), "CTXT")
self.assertEqual(dns.rdatatype.from_text("CTXT"), CTXT)
self.assertEqual(dns.rdatatype.RdataType.make("CTXT"), CTXT)
self.assertEqual(dns.rdatatype.from_text("ctxt"), CTXT)
self.assertEqual(dns.rdatatype.RdataType.make("ctxt"), CTXT)

def test_module_reregistration(self):
def bad():
TTXTTWO = dns.rdatatype.TXT
Expand Down

0 comments on commit f972afd

Please sign in to comment.