Add arvados.getjobparam() convenience function.

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

diff --git a/sdk/python/arvados.py b/sdk/python/arvados.py
index c1890c4da2058366523c66691fca7fe13b3a320f..f4fe0707865c22e37cd6e4eab317d31723d13372 100644 (file)
--- a/sdk/python/arvados.py
+++ b/sdk/python/arvados.py
@@ -93,6 +93,9 @@ def current_job():
     _current_job = t
     return t
 
+def getjobparam(*args):
+    return current_job()['script_parameters'].get(*args)
+
 def api():
     return service
 
@@ -149,16 +152,29 @@ 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):
-        if 'stdin' not in kwargs:
-            kwargs['stdin'] = subprocess.PIPE
-        if 'stdout' not in kwargs:
-            kwargs['stdout'] = subprocess.PIPE
-        if 'stderr' not in kwargs:
-            kwargs['stderr'] = subprocess.PIPE
-        p = subprocess.Popen(execargs, close_fds=True, shell=False,
-                             **kwargs)
+        kwargs.setdefault('stdin', subprocess.PIPE)
+        kwargs.setdefault('stdout', subprocess.PIPE)
+        kwargs.setdefault('stderr', sys.stderr)
+        kwargs.setdefault('close_fds', True)
+        kwargs.setdefault('shell', False)
+        p = subprocess.Popen(execargs, **kwargs)
         stdoutdata, stderrdata = p.communicate(None)
         if p.returncode != 0:
             raise Exception("run_command %s exit %d:\n%s" %
@@ -322,6 +338,11 @@ class util:
         collection -- collection locator
         path -- where to extract: absolute, or relative to job tmp
         """
+        matches = re.search(r'^([0-9a-f]+)(\+[\w@]+)*$', collection)
+        if matches:
+            collection_hash = matches.group(1)
+        else:
+            collection_hash = hashlib.md5(collection).hexdigest()
         if not re.search('^/', path):
             path = os.path.join(current_job().tmpdir, path)
         lockfile = open(path + '.lock', 'w')
@@ -332,7 +353,7 @@ class util:
             os.mkdir(path)
         already_have_it = False
         try:
-            if os.readlink(os.path.join(path, '.locator')) == collection:
+            if os.readlink(os.path.join(path, '.locator')) == collection_hash:
                 already_have_it = True
         except OSError:
             pass
@@ -364,7 +385,7 @@ class util:
                     outfile.close()
         if len(files_got) < len(files):
             raise Exception("Wanted files %s but only got %s from %s" % (files, files_got, map(lambda z: z.name(), list(CollectionReader(collection).all_files()))))
-        os.symlink(collection, os.path.join(path, '.locator'))
+        os.symlink(collection_hash, os.path.join(path, '.locator'))
 
         lockfile.close()
         return path
@@ -674,7 +695,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:
@@ -695,6 +716,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)
@@ -731,7 +756,7 @@ class CollectionWriter(object):
     def set_current_stream_name(self, newstreamname):
         if re.search(r'[ \t\n]', newstreamname):
             raise AssertionError("Manifest stream names cannot contain whitespace")
-        self._current_stream_name = newstreamname
+        self._current_stream_name = '.' if newstreamname=='' else newstreamname
     def current_stream_name(self):
         return self._current_stream_name
     def finish_current_stream(self):