From c508257d6478e75a74a2fe0c9e770ef89c50eb41 Mon Sep 17 00:00:00 2001
From: Ben Greiner <code@bnavigator.de>
Date: Mon, 18 Jan 2021 19:05:03 +0100
Subject: [PATCH] add support for dnspython 2.1

---
 localzone/context.py  | 69 +++++++++++++++++++++++++++++--------------
 tests/test_context.py |  5 +++-
 tests/test_models.py  |  3 +-
 3 files changed, 53 insertions(+), 24 deletions(-)

diff --git a/localzone/context.py b/localzone/context.py
index 8da5048..7354c4a 100644
--- a/localzone/context.py
+++ b/localzone/context.py
@@ -13,6 +13,14 @@
 import dns.rdataclass
 import dns.tokenizer
 import dns.zone
+try:
+    # The api for the zonefile reader was exposed in dnspython 2.1
+    from dns.zonefile import Reader
+    DNSPYTHON21 = True
+except ImportError:
+    from dns.zone import _MasterReader
+    DNSPYTHON21 = False
+
 from .models import Zone
 
 
@@ -57,25 +65,42 @@ def load(filename, origin=None):
     """
     with open(filename) as text:
         tok = dns.tokenizer.Tokenizer(text, filename)
-        reader = dns.zone._MasterReader(
-            tok,
-            origin=origin,
-            rdclass=dns.rdataclass.IN,
-            relativize=True,
-            zone_factory=Zone,
-            allow_include=True,
-            check_origin=True,
-        )
-        reader.read()
-
-        # TODO: remember that any method using the zone.filename property should
-        # have it passed as a parameter
-        reader.zone._filename = filename
-
-        # starting with dnspython v1.16.0, use default_ttl
-        try:
-            reader.zone._ttl = reader.default_ttl
-        except AttributeError:
-            reader.zone._ttl = reader.ttl
-
-        return reader.zone
+        if DNSPYTHON21:
+            zone = Zone(
+                origin,
+                dns.rdataclass.IN,
+                relativize=True,
+                )
+            with zone.writer() as txn:
+                reader = Reader(
+                    tok,
+                    rdclass=dns.rdataclass.IN,
+                    txn=txn,
+                    allow_include=True,
+                )
+                reader.read()
+        else:
+            reader = _MasterReader(
+                tok,
+                origin=origin,
+                rdclass=dns.rdataclass.IN,
+                relativize=True,
+                zone_factory=Zone,
+                allow_include=True,
+                check_origin=True,
+            )
+            reader.read()
+            zone = reader.zone
+
+    # TODO: remember that any method using the zone.filename property should
+    # have it passed as a parameter
+    zone._filename = filename
+
+    # starting with dnspython v1.16.0, use default_ttl
+    try:
+        zone._ttl = reader.default_ttl
+    except AttributeError:
+        zone._ttl = reader.ttl
+
+    return zone
+
diff --git a/tests/test_context.py b/tests/test_context.py
index 2a32d00..baac084 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -1,6 +1,9 @@
 import pytest
 import localzone
-from dns.zone import UnknownOrigin
+try:
+    from dns.zonefile import UnknownOrigin
+except ImportError:
+    from dns.zone import UnknownOrigin
 
 ZONEFILE = "tests/zonefiles/db.example.com"
 ORIGIN = "example.com."
diff --git a/tests/test_models.py b/tests/test_models.py
index 9ba0c40..621f6c0 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1,4 +1,5 @@
 from dns.rdatatype import UnknownRdatatype
+from dns.exception import SyntaxError as DNSSyntaxError
 import pytest
 import localzone
 
@@ -77,7 +78,7 @@ def test_zone_add_record_unknown_type():
 
 def test_zone_add_record_no_content():
     with localzone.manage(ZONEFILE, ORIGIN) as z:
-        with pytest.raises(AttributeError):
+        with pytest.raises((AttributeError, DNSSyntaxError)):
             z.add_record("test", "txt", None)