Accept data from a generator in arvados.CollectionWriter.write()

[arvados.git] / sdk / python / arvados.py

diff --git a/sdk/python/arvados.py b/sdk/python/arvados.py
index 9c92f7a6c255ba0efca06162078e94310b82d9ba..354ef629f8641d27f20935c4b57de0df5808520c 100644 (file)
--- a/sdk/python/arvados.py
+++ b/sdk/python/arvados.py
@@ -149,11 +149,26 @@ class job_setup:
             exit(0)
 
 class util:
+    @staticmethod
+    def clear_tmpdir(path=None):
+        """
+        Ensure the given directory (or TASK_TMPDIR if none given)
+        exists and is empty.
+        """
+        if path == None:
+            path = current_task().tmpdir
+        if os.path.exists(path):
+            p = subprocess.Popen(['rm', '-rf', path])
+            stdout, stderr = p.communicate(None)
+            if p.returncode != 0:
+                raise Exception('rm -rf %s: %s' % (path, stderr))
+        os.mkdir(path)
+
     @staticmethod
     def run_command(execargs, **kwargs):
         kwargs.setdefault('stdin', subprocess.PIPE)
         kwargs.setdefault('stdout', subprocess.PIPE)
-        kwargs.setdefault('stderr', subprocess.PIPE)
+        kwargs.setdefault('stderr', sys.stderr)
         kwargs.setdefault('close_fds', True)
         kwargs.setdefault('shell', False)
         p = subprocess.Popen(execargs, **kwargs)
@@ -677,7 +692,7 @@ class CollectionWriter(object):
         self.start_new_stream(stream_name)
         todo = []
         if max_manifest_depth == 0:
-            dirents = util.listdir_recursive(path)
+            dirents = sorted(util.listdir_recursive(path))
         else:
             dirents = sorted(os.listdir(path))
         for dirent in dirents:
@@ -698,6 +713,10 @@ class CollectionWriter(object):
         map(lambda x: self.write_directory_tree(*x), todo)
 
     def write(self, newdata):
+        if hasattr(newdata, '__iter__'):
+            for s in newdata:
+                self.write(s)
+            return
         self._data_buffer += [newdata]
         self._data_buffer_len += len(newdata)
         self._current_stream_length += len(newdata)