ENH: support of .gz in create_tree and ok_file_has_content + minor RF

datalad/tests/test_utils.py

@@ -63,6 +63,7 @@
 from ..utils import map_items
 from ..utils import unlink
 from ..utils import CMD_MAX_ARG
+from ..utils import create_tree
 from ..support.annexrepo import AnnexRepo
 
 from nose.tools import (
@@ -91,7 +92,8 @@
 from .utils import skip_if_no_module
 from .utils import (
     probe_known_failure, skip_known_failure, known_failure, known_failure_v6,
-    known_failure_direct_mode, skip_if
+    known_failure_direct_mode, skip_if,
+    ok_file_has_content
 )
 
 
@@ -1186,4 +1188,22 @@ def test_CMD_MAX_ARG():
     # 100 is arbitrarily large small integer ;)
     # if fails -- we are unlikely to be able to work on this system
     # and something went really wrong!
-    assert_greater(CMD_MAX_ARG, 100)
+    assert_greater(CMD_MAX_ARG, 100)
+
+
+@with_tempfile(mkdir=True)
+def test_create_tree(path):
+    content = u"мама мыла раму"
+    create_tree(path, OrderedDict([
+        ('1', content),
+        ('sd', OrderedDict(
+            [
+            # right away an obscure case where we have both 1 and 1.gz
+                ('1', content*2),
+                ('1.gz', content*3),
+            ]
+        )),
+    ]))
+    ok_file_has_content(op.join(path, '1'), content)
+    ok_file_has_content(op.join(path, 'sd', '1'), content*2)
+    ok_file_has_content(op.join(path, 'sd', '1.gz'), content*3, decompress=True)

datalad/tests/utils.py

@@ -9,6 +9,7 @@
 """Miscellaneous utilities to assist with testing"""
 
 import glob
+import gzip
 import inspect
 import shutil
 import stat
@@ -391,19 +392,30 @@ def ok_exists(path):
     assert exists(path), 'path %s does not exist' % path
 
 
-def ok_file_has_content(path, content, strip=False, re_=False, **kwargs):
+def ok_file_has_content(path, content, strip=False, re_=False,
+                        decompress=False, **kwargs):
     """Verify that file exists and has expected content"""
     ok_exists(path)
-    with open(path, 'r') as f:
-        content_ = f.read()
+    if decompress:
+        if path.endswith('.gz'):
+            open_func = gzip.open
+        else:
+            raise NotImplementedError("Don't know how to decompress %s" % path)
+    else:
+        open_func = open
+
+    with open_func(path, 'rb') as f:
+        file_content = f.read()
+        if isinstance(content, text_type):
+            file_content = assure_unicode(file_content)
 
         if strip:
-            content_ = content_.strip()
+            file_content = file_content.strip()
 
         if re_:
-            assert_re_in(content, content_, **kwargs)
+            assert_re_in(content, file_content, **kwargs)
         else:
-            assert_equal(content, content_, **kwargs)
+            assert_equal(content, file_content, **kwargs)
 
 
 #

datalad/utils.py

@@ -21,6 +21,7 @@
 import platform
 import gc
 import glob
+import gzip
 import string
 import wrapt
 
@@ -2132,15 +2133,11 @@ def create_tree(path, tree, archives_leading_dir=True, remove_existing=False):
                     archives_leading_dir=archives_leading_dir,
                     remove_existing=remove_existing)
         else:
-            if PY2:
-                open_kwargs = {'mode': "w"}
-                if isinstance(load, text_type):
-                    load = load.encode('utf-8')
-            else:
-                open_kwargs = {'mode': "w", 'encoding': "utf-8"}
-
-            with open(full_name, **open_kwargs) as f:
-                f.write(load)
+            open_func = open
+            if full_name.endswith('.gz'):
+                open_func = gzip.open
+            with open_func(full_name, "wb") as f:
+                f.write(assure_bytes(load, 'utf-8'))
         if executable:
             os.chmod(full_name, os.stat(full_name).st_mode | stat.S_IEXEC)