Merge pull request #7 from toddw-as/patch-1

Allow an optional file encoding (codec).

atomicfile.py

@@ -3,6 +3,7 @@
 import errno
 import os
 import tempfile
+import codecs
 
 
 umask = os.umask(0)
@@ -46,11 +47,16 @@ class AtomicFile(object):
     the temporary copy to the original name, making the changes visible.
     If the object is destroyed without being closed, all your writes are
     discarded.
+    If an ``encoding`` argument is specified, codecs.open will be called to open
+    the file in the wanted encoding.
     """
-    def __init__(self, name, mode="w+b", createmode=None):
+    def __init__(self, name, mode="w+b", createmode=None, encoding=None):
         self.__name = name  # permanent name
         self._tempname = _maketemp(name, createmode=createmode)
-        self._fp = open(self._tempname, mode)
+        if encoding:
+            self._fp = codecs.open(self._tempname, mode, encoding)
+        else:
+            self._fp = open(self._tempname, mode)
 
         # delegated methods
         self.write = self._fp.write

tests.py

@@ -2,6 +2,7 @@
 
 import os
 import unittest
+import codecs
 
 from atomicfile import AtomicFile
 
@@ -73,6 +74,26 @@ def test_permissions(self):
         finally:
             os.remove(self.filename)
 
+    def test_encoding(self):
+        data = u"Unicode Capit\xe1n is written by AtomicFile.\n"
+        encoding = "utf-8"
+        af = AtomicFile(self.filename, "wb", encoding=encoding)
+        af.write(data)
+        af.close()
+
+        f = codecs.open(self.filename, "rb", encoding=encoding)
+        decoded_result = f.read()
+        f.close()
+        f = open(self.filename, "rb")
+        raw_result = f.read()
+        f.close()
+
+        try:
+            self.assertEqual(data, decoded_result)
+            self.assertEqual(data.encode(encoding), raw_result)
+        finally:
+            os.remove(self.filename)
+
 
 if __name__ == "__main__":
     unittest.main()