rework the write file and chunk helpers to do intelligent 'mode' selection based on input data type (unicode or bytes) as a follow-up to PR48. Thoroughly exercise the helpers using all four combinations of format and mode in unit tests to make sure we cover the cases previously handled by explicit byte and unicode conversion with force utf8 and force native str helpers when running in python3 on 'experimental' branch.

git-svn-id: svn+ssh://svn.code.sf.net/p/migrid/code/trunk@6040 b75ad72c-e7d7-11dd-a971-7dbc132099af

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()