Accommodating pipes in programs.py

Adding in docker pipes unit test

Author: cmarkello

src/toil_scripts/lib/programs.py

@@ -25,11 +25,13 @@ def docker_call(tool,
                 outputs=None,
                 docker_parameters=None,
                 check_output=False,
+                return_stderr=False,
                 mock=None):
     """
     Calls Docker, passing along parameters and tool.
 
-    :param str tool: Name of the Docker image to be used (e.g. quay.io/ucsc_cgl/samtools)
+    :param str OR list[str] tool: Single str or list of names of the Docker images and order to be used 
+                      (e.g. quay.io/ucsc_cgl/samtools)
     :param list[str] parameters: Command line arguments to be passed to the tool
     :param str work_dir: Directory to mount into the container via `-v`. Destination convention is /data
     :param bool rm: Set to True to pass `--rm` flag.
@@ -41,8 +43,24 @@ def docker_call(tool,
                                   or a url. The value is only used if mock=True
     :param dict[str,str] docker_parameters: Parameters to pass to docker
     :param bool check_output: When True, this function returns docker's output
+    :param bool return_stderr: When True, this function includes stderr in docker's output
     :param bool mock: Whether to run in mock mode. If this variable is unset, its value will be determined by
                       the environment variable.
+    
+    Piping docker command examples:
+        Running 'pipe-in-single-container' mode for command
+        'head -c 1M </dev/urandom | tee >(md5sum 1>&2) | gzip | gunzip | md5sum 1>&2':
+            command= ['head -c 1M </dev/urandom | tee >(md5sum 1>&2)', 'gzip -', 'gunzip -', 'md5sum - 1>&2']
+            work_dir=curr_work_dir
+            docker_tools=['ubuntu']
+            stdout = docker_call(work_dir=docker_work_dir, parameters=command, tool=docker_tools, check_output=True)
+        
+        Running 'pipe-of-containers' mode for command
+        'head -c 1M </dev/urandom | tee >(md5sum 1>&2) | gzip | gunzip | md5sum 1>&2':
+            command= ['head -c 1M </dev/urandom | tee >(md5sum 1>&2)', 'gzip -', 'gunzip -', 'md5sum - 1>&2']
+            work_dir=curr_work_dir
+            docker_tools=['ubuntu', 'ubuntu', 'ubuntu', 'ubuntu']
+            stdout = docker_call(work_dir=docker_work_dir, parameters=command, tool=docker_tools, check_output=True)
     """
     from toil_scripts.lib.urls import download_url
 
@@ -83,21 +101,53 @@ def docker_call(tool,
     if env:
         for e, v in env.iteritems():
             base_docker_call.extend(['-e', '{}={}'.format(e, v)])
+    
     if docker_parameters:
         base_docker_call += docker_parameters
-
-    _log.debug("Calling docker with %s." % " ".join(base_docker_call + [tool] + parameters))
-
-    docker_call = base_docker_call + [tool] + parameters
-
+   
+    docker_call = []
+    
+    # Pipe functionality
+    #   each element in the parameters list must represent a sub-pipe command
+    shell_flag = False      # Flag for running subprocess with string command or list command
+    if isinstance(tool, list):
+        command_list = []
+        if len(tool) > 1:
+            # If tool is a list containing multiple docker container name strings
+            #   then format the docker call in the 'pipe-of-containers' mode
+            docker_call.extend(base_docker_call + ['--entrypoint /bin/bash', tool[0], '-c \'{}\''.format(parameters[0])])
+            for i in xrange(1, len(tool)):
+                docker_call.extend(['|'] + base_docker_call + ['-i --entrypoint /bin/bash', tool[i], '-c \'{}\''.format(parameters[i])])
+            docker_call = " ".join(docker_call)
+            shell_flag=True
+            print("Calling docker with {}".format(docker_call)) 
+            _log.debug("Calling docker with %s." % docker_call)
+            
+        elif len(tool) == 1:
+            # If tool is a list containing a single docker container name string
+            #   then format the docker call in the 'pipe-in-single-container' mode
+            docker_call.extend(base_docker_call + ['--entrypoint /bin/bash', tool[0], '-c \'{}\''.format(" | ".join(parameters))])
+            docker_call = " ".join(docker_call)
+            shell_flag=True
+            print("Calling docker with {}".format(docker_call)) 
+            _log.debug("Calling docker with %s." % docker_call)
+             
+    else:        
+        docker_call = base_docker_call + [tool] + parameters
+        _log.debug("Calling docker with %s." % " ".join(docker_call))
+    
     try:
         if outfile:
-            subprocess.check_call(docker_call, stdout=outfile)
+            subprocess.check_call(docker_call, stdout=outfile, shell=shell_flag)
         else:
             if check_output:
-                return subprocess.check_output(docker_call)
+                #TODO change stderr
+                if return_stderr:
+                    return subprocess.check_output(docker_call, shell=shell_flag, stderr=subprocess.STDOUT)
+                else:
+                    return subprocess.check_output(docker_call, shell=shell_flag)
             else:
-                subprocess.check_call(docker_call)
+                subprocess.check_call(docker_call, shell=shell_flag)
     # Fix root ownership of output files
     except:
         # Panic avoids hiding the exception raised in the try block
@@ -122,5 +172,19 @@ def _fix_permissions(base_docker_call, tool, work_dir):
     """
     base_docker_call.append('--entrypoint=chown')
     stat = os.stat(work_dir)
-    command = base_docker_call + [tool] + ['-R', '{}:{}'.format(stat.st_uid, stat.st_gid), '/data']
-    subprocess.check_call(command)
+    command = []
+    command_list = []
+    if isinstance(tool, list):
+        for i in xrange(len(tool)):
+            command = base_docker_call + [tool[i]] + ['-R', '{}:{}'.format(stat.st_uid, stat.st_gid), '/data']
+            command_list.append(command)
+    else:
+        command = base_docker_call + [tool] + ['-R', '{}:{}'.format(stat.st_uid, stat.st_gid), '/data']
+        command_list.append(command)
+    
+    for i in xrange(len(command_list)):
+        subprocess.check_call(command_list[i])
+
+
+
+

src/toil_scripts/lib/test/test_programs.py

@@ -1,5 +1,5 @@
 import os
-
+import re
 
 def test_docker_call(tmpdir):
     from toil_scripts.lib.programs import docker_call
@@ -12,3 +12,18 @@ def test_docker_call(tmpdir):
     with open(fpath, 'w') as f:
         docker_call(tool='ubuntu', env=dict(foo='bar'), parameters=['printenv', 'foo'], outfile=f)
     assert open(fpath).read() == 'bar\n'
+
+    # Test pipe functionality
+    # download ubuntu docker image
+    docker_call(work_dir=work_dir, tool="ubuntu")
+    command = ['head -c 1G </dev/urandom | tee >(md5sum 1>&2)', 'gzip -', 'gunzip -', 'md5sum - 1>&2']
+    # Test 'pipe-in-single-container' mode
+    docker_tools1=['ubuntu']
+    stdout1 = docker_call(work_dir=work_dir, parameters=command, tool=docker_tools1, check_output=True, return_stderr=True)
+    test1 = re.findall(r"([a-fA-F\d]{32})", stdout1)
+    assert test1[0] == test1[1]
+    # Test 'pipe-of-containers' mode
+    docker_tools2=['ubuntu', 'ubuntu', 'ubuntu', 'ubuntu']
+    stdout2 = docker_call(work_dir=work_dir, parameters=command, tool=docker_tools2, check_output=True, return_stderr=True)
+    test2 = re.findall(r"([a-fA-F\d]{32})", stdout2)
+    assert test2[0] == test2[1]