Merge remote-tracking branch 'origin/master' into edge

Author: jonasbardino

mig/shared/fileio.py

@@ -70,6 +70,21 @@
     exit(1)
 
 
+def _auto_adjust_mode(data, mode):
+    """Select suitable file open mode based on string type of data. I.e. whether
+    to use binary or text mode depending on data in bytes or unicode format.
+    """
+
+    # NOTE: detect byte/unicode writes and handle explicitly in a portable way
+    if isinstance(data, bytes):
+        if 'b' not in mode:
+            mode = "%sb" % mode  # appended to avoid mode ordering error on PY2
+    else:
+        if 'b' in mode:
+            mode = mode.strip('b')
+    return mode
+
+
 def _write_chunk(path, chunk, offset, logger=None, mode='r+b',
                  make_parent=True, create_file=True):
     """Internal helper to wrap writing of chunks with offset to path.
@@ -104,8 +119,8 @@ def _write_chunk(path, chunk, offset, logger=None, mode='r+b',
 
     try:
         with open(path, mode) as filehandle:
+            # Make sure we can write at requested position, filling if needed
             if offset > 0:
-                # Make sure we can write at requested position, filling if needed
                 try:
                     filehandle.seek(offset)
                 except:
@@ -131,9 +146,7 @@ def write_chunk(path, chunk, offset, logger, mode='r+b'):
     if not logger:
         logger = null_logger("dummy")
 
-    # NOTE: detect byte writes and handle explicitly in a portable way
-    if isinstance(chunk, bytes) and 'b' not in mode:
-        mode = "%sb" % mode  # appended to avoid mode ordering error on PY2
+    mode = _auto_adjust_mode(chunk, mode)
 
     return _write_chunk(path, chunk, offset, logger, mode)
 
@@ -147,9 +160,7 @@ def write_file(content, path, logger, mode='w', make_parent=True, umask=None):
     if umask is not None:
         old_umask = os.umask(umask)
 
-    # NOTE: detect byte writes and handle explicitly in a portable way
-    if isinstance(content, bytes) and 'b' not in mode:
-        mode = "%sb" % mode  # appended to avoid mode ordering error on PY2
+    mode = _auto_adjust_mode(content, mode)
 
     retval = _write_chunk(path, content, offset=0, logger=logger, mode=mode,
                           make_parent=make_parent, create_file=False)

tests/test_mig_shared_fileio.py

@@ -12,6 +12,8 @@
 
 DUMMY_BYTES = binascii.unhexlify('DEADBEEF')  # 4 bytes
 DUMMY_BYTES_LENGTH = 4
+DUMMY_UNICODE = u'UniCode123'
+DUMMY_UNICODE_LENGTH = len(DUMMY_UNICODE)
 DUMMY_FILE_WRITECHUNK = 'fileio/write_chunk'
 DUMMY_FILE_WRITEFILE = 'fileio/write_file'
 
@@ -65,6 +67,33 @@ def test_store_bytes_at_offset(self):
                              "expected a hole was left")
             self.assertEqual(content[3:], DUMMY_BYTES)
 
+    def test_store_bytes_in_text_mode(self):
+        fileio.write_chunk(self.tmp_path, DUMMY_BYTES, 0, self.logger,
+                           mode="r+")
+
+        with open(self.tmp_path, 'rb') as file:
+            content = file.read(1024)
+            self.assertEqual(len(content), DUMMY_BYTES_LENGTH)
+            self.assertEqual(content[:], DUMMY_BYTES)
+
+    def test_store_unicode(self):
+        fileio.write_chunk(self.tmp_path, DUMMY_UNICODE, 0, self.logger,
+                           mode='r+')
+
+        with open(self.tmp_path, 'r') as file:
+            content = file.read(1024)
+            self.assertEqual(len(content), DUMMY_UNICODE_LENGTH)
+            self.assertEqual(content[:], DUMMY_UNICODE)
+
+    def test_store_unicode_in_binary_mode(self):
+        fileio.write_chunk(self.tmp_path, DUMMY_UNICODE, 0, self.logger,
+                           mode='r+b')
+
+        with open(self.tmp_path, 'r') as file:
+            content = file.read(1024)
+            self.assertEqual(len(content), DUMMY_UNICODE_LENGTH)
+            self.assertEqual(content[:], DUMMY_UNICODE)
+
 
 class MigSharedFileio__write_file(MigTestCase):
     def setUp(self):
@@ -103,6 +132,36 @@ def test_store_bytes(self):
             self.assertEqual(len(content), DUMMY_BYTES_LENGTH)
             self.assertEqual(content[:], DUMMY_BYTES)
 
+    def test_store_bytes_in_text_mode(self):
+        did_succeed = fileio.write_file(DUMMY_BYTES, self.tmp_path, self.logger,
+                                        mode="w")
+        self.assertTrue(did_succeed)
+
+        with open(self.tmp_path, 'rb') as file:
+            content = file.read(1024)
+            self.assertEqual(len(content), DUMMY_BYTES_LENGTH)
+            self.assertEqual(content[:], DUMMY_BYTES)
+
+    def test_store_unicode(self):
+        did_succeed = fileio.write_file(DUMMY_UNICODE, self.tmp_path,
+                                        self.logger, mode='w')
+        self.assertTrue(did_succeed)
+
+        with open(self.tmp_path, 'r') as file:
+            content = file.read(1024)
+            self.assertEqual(len(content), DUMMY_UNICODE_LENGTH)
+            self.assertEqual(content[:], DUMMY_UNICODE)
+
+    def test_store_unicode_in_binary_mode(self):
+        did_succeed = fileio.write_file(DUMMY_UNICODE, self.tmp_path,
+                                        self.logger, mode='wb')
+        self.assertTrue(did_succeed)
+
+        with open(self.tmp_path, 'r') as file:
+            content = file.read(1024)
+            self.assertEqual(len(content), DUMMY_UNICODE_LENGTH)
+            self.assertEqual(content[:], DUMMY_UNICODE)
+
 
 if __name__ == '__main__':
     testmain()