diff --git a/atomicfile.py b/atomicfile.py index 596dff5..37e92d8 100644 --- a/atomicfile.py +++ b/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 diff --git a/tests.py b/tests.py index e3f80c5..aa2e858 100644 --- a/tests.py +++ b/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()