Merge pull request #97 from encukou/zone-to_file

Fix failing test for Zone.to_file

2 files changed, 30 insertions, 9 deletions

diff --git a/dns/zone.py b/dns/zone.py
index 9193a7d..b162ad3 100644
--- a/dns/zone.py
+++ b/dns/zone.py
@@ -21,6 +21,7 @@ import builtins
 import io
 import re
 import sys
+import os
 
 import dns.exception
 import dns.name
@@ -444,7 +445,7 @@ class Zone(object):
                     for rdata in rds:
                         yield (name, rds.ttl, rdata)
 
-    def to_file(self, f, sorted=True, relativize=True, nl=None):
+    def to_file(self, f, sorted=True, relativize=True, nl=None, binary=False):
         """Write a zone to a file.
 
         @param f: file or string.  If I{f} is a string, it is treated
@@ -460,17 +461,23 @@ class Zone(object):
         output will use the platform's native end-of-line marker (i.e.
         LF on POSIX, CRLF on Windows, CR on Macintosh).
         @type nl: string or None
+        @param binary: True if the file is open in binary mode
+        @type binary: bool
         """
 
         if nl is None:
-            opts = 'w'
-        else:
-            opts = 'wb'
+            nl = os.linesep
+
         if isinstance(f, str):
-            f = open(f, opts)
+            if binary:
+                mode = 'wb'
+            else:
+                mode = 'w'
+            f = open(f, mode)
             want_close = True
         else:
             want_close = False
+
         try:
             if sorted:
                 names = builtins.sorted(self.keys())
@@ -479,11 +486,12 @@ class Zone(object):
             for n in names:
                 l = self[n].to_text(n, origin=self.origin,
                                     relativize=relativize)
-                if nl is None:
-                    print(l, file=f)
-                else:
+                if binary:
                     f.write(l.encode('ascii'))
                     f.write(nl.encode('ascii'))
+                else:
+                    f.write(l)
+                    f.write(nl)
         finally:
             if want_close:
                 f.close()
diff --git a/tests/test_zone.py b/tests/test_zone.py
index 67fc2ae..08557ad 100644
--- a/tests/test_zone.py
+++ b/tests/test_zone.py
@@ -121,12 +121,25 @@ class ZoneTestCase(unittest.TestCase):
                 os.unlink('example2.out')
         self.assertTrue(ok)
 
+    def testFromFile2b(self):
+        """Test to_file with a binary file"""
+        z = dns.zone.from_file('example', 'example', relativize=False)
+        ok = False
+        try:
+            with open('example2b.out', 'wb') as f:
+                z.to_file(f, relativize=False, nl='\x0a', binary=True)
+            ok = filecmp.cmp('example2b.out', 'example2.good')
+        finally:
+            if not _keep_output:
+                os.unlink('example2b.out')
+        self.assertTrue(ok)
+
     def testToText(self):
         z = dns.zone.from_file('example', 'example')
         ok = False
         try:
             text_zone = z.to_text(nl='\x0a')
-            f = open('example3.out', 'wb')
+            f = open('example3.out', 'w')
             f.write(text_zone)
             f.close()
             ok = filecmp.cmp('example3.out', 'example3.good')