From 4f23be59c70c552f9ec4f0f9cbbe70d8aec01cb5 Mon Sep 17 00:00:00 2001
From: Konrad Karczewski <konradjkarczewski@gmail.com>
Date: Mon, 11 Feb 2019 11:50:58 -0500
Subject: [PATCH] Add python files or folders from environment variable; zip
 files together (#127)

* Add python files or folders from environment variable; zip files together

* bumping version

* files -> pyfiles

* missed one

* overloaded variable
---
 cloudtools/submit.py | 34 +++++++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/cloudtools/submit.py b/cloudtools/submit.py
index c8896cc2a79..82dbf9fcea5 100644
--- a/cloudtools/submit.py
+++ b/cloudtools/submit.py
@@ -1,9 +1,19 @@
 from subprocess import check_call
+import os
+import tempfile
+import zipfile
+
+try:
+    standard_scripts = os.environ['HAIL_SCRIPTS'].split(':')
+except Exception:
+    standard_scripts = None
+
 
 def init_parser(parser):
     parser.add_argument('name', type=str, help='Cluster name.')
     parser.add_argument('script', type=str)
     parser.add_argument('--files', required=False, type=str, help='Comma-separated list of files to add to the working directory of the Hail application.')
+    parser.add_argument('--pyfiles', required=False, type=str, help='Comma-separated list of files (or directories with python files) to add to the PYTHONPATH.')
     parser.add_argument('--properties', '-p', required=False, type=str, help='Extra Spark properties to set.')
     parser.add_argument('--args', type=str, help='Quoted string of arguments to pass to the Hail script being submitted.')
 
@@ -15,6 +25,28 @@ def main(args):
     files = ''
     if args.files:
         files = args.files
+    pyfiles = []
+    if args.pyfiles:
+        pyfiles.extend(args.files.split(','))
+    if standard_scripts:
+        pyfiles.extend(standard_scripts)
+    if pyfiles:
+        tfile = tempfile.mkstemp(suffix='.zip', prefix='pyscripts_')[1]
+        zipf = zipfile.ZipFile(tfile, 'w', zipfile.ZIP_DEFLATED)
+        for hail_script_entry in pyfiles:
+            if hail_script_entry.endswith('.py'):
+                zipf.write(hail_script_entry, arcname=os.path.basename(hail_script_entry))
+            else:
+                for root, _, pyfiles_walk in os.walk(hail_script_entry):
+                    for pyfile in pyfiles_walk:
+                        if pyfile.endswith('.py'):
+                            zipf.write(os.path.join(root, pyfile),
+                                       os.path.relpath(os.path.join(root, pyfile),
+                                                       os.path.join(hail_script_entry, '..')))
+        zipf.close()
+        pyfiles = tfile
+    else:
+        pyfiles = ''
 
     # create properties argument
     properties = ''
@@ -31,7 +63,7 @@ def main(args):
         args.script,
         '--cluster={}'.format(args.name),
         '--files={}'.format(files),
-        # '--py-files={}'.format(zip_path),
+        '--py-files={}'.format(pyfiles),
         '--properties={}'.format(properties)
     ]