Various fixes

2 files changed, 89 insertions, 23 deletions

diff --git a/condor_split.py b/condor_split.py
index faf08d8..480c94d 100755
--- a/condor_split.py
+++ b/condor_split.py
@@ -3,10 +3,15 @@ import argparse
 import os
 import sys
 import fnmatch
+import shutil
+import tempfile
 
-def path_dive(p, ext, strip_components):
+VERBOSE = False
+
+def path_search(p, ext, strip_components):
     filenames = []
     p = os.path.abspath(p)
+    index = 1
     for root, _, files in os.walk(p):
         for f in files:
             path = os.path.join(root, f)
@@ -20,29 +25,91 @@ def path_dive(p, ext, strip_components):
                         components = head.split(os.path.sep)[strip_components:]
                         path = os.path.join(os.path.sep.join(components), tail)
                     else:
-                        #strip_components = strip_max
+                        # strip_components = strip_max
                         print("Warning: Cannot strip {} components from {} (max {})".format(strip_components, path, strip_max))
+
+                if VERBOSE:
+                    print("[{}]:{}".format(index, path))
                 filenames.append(path)
+                index += 1
     return filenames
 
-def assign_chunk():
-    pass
+def generate_manifest(manifest):
+    dest = tempfile.NamedTemporaryFile("w+", delete=False)
+    for line in manifest:
+        dest.write(line + os.linesep)
+    return dest.name
+
+def assign_chunks(manifest, chunks, output_dir, as_jobs=False):
+    label = 0
+    count = 0
+    written = 0
+    # max_chunks = 0
+    dest = None
+    output_dir = os.path.abspath(output_dir)
+
+    if VERBOSE:
+        print("Output directory: {}".format(output_dir))
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    else:
+        shutil.rmtree(output_dir)
+        os.makedirs(output_dir)
+
+    if as_jobs:
+        manifest_count = len(file(manifest, 'r').readlines())
+        max_chunks = manifest_count / chunks
+        chunks = max_chunks
+
+
+    for line in open(manifest, 'r'):
+        line = line.strip(os.linesep)
+        line += " "
+        if count % chunks == 0:
+            if dest:
+                dest.close()
+                if VERBOSE:
+                    print("[{}]: {}, {} entries".format(label, os.path.basename(dest.name), written))
+                written = 0
+            dest = file(os.path.join(output_dir, "stdin." + str(label)), 'w')
+            label += 1
+        dest.write(line)
+        count += 1
+        written += 1
+    os.unlink(manifest)
+
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("search_path", action="store", help="Directory to search under")
     parser.add_argument("--strip-components", default=0, type=int, help="Top-level directories to strip")
     parser.add_argument("--output-dir", default="input")
-    parser.add_argument("--extension", default="*.*", help="Extension of files")
-    parser.add_argument("--chunks", default=4, help="Number of entries per file")
+    parser.add_argument("--extension", default="*.*", type=str, help="Extension of files")
+    parser.add_argument("--chunks", default=1, type=int, help="Number of entries per file")
+    parser.add_argument("--as-jobs", action="store_true", help="Number of expected jobs")
+    parser.add_argument("--verbose", action="store_true")
+    parser.add_argument("search_path", action="store", type=str, help="Directory to search under")
     args = parser.parse_args()
 
-    paths = path_dive(args.search_path, args.extension, args.strip_components)
-    print("{} files".format(len(paths)))
-    if not paths:
-        exit(0)
+    global VERBOSE
+    VERBOSE = args.verbose
+    chunks = args.chunks
+    if chunks < 1:
+        chunks = 1
+    index_at = 0
 
+    files = path_search(args.search_path, args.extension, args.strip_components)
+    if files:
+        index_at = 1
+
+    if VERBOSE:
+        print("{} files".format(len(files)))
+
+    if not files:
+        exit(0)
 
+    manifest = generate_manifest(files)
+    assign_chunks(manifest, chunks, args.output_dir, as_jobs=args.as_jobs)
 
 if __name__ == "__main__":
     main()
diff --git a/condor_wrap.py b/condor_wrap.py
index 2142c88..0aa25c4 100755
--- a/condor_wrap.py
+++ b/condor_wrap.py
@@ -15,21 +15,20 @@ def main():
 
     job = [os.path.normpath(args.job)]
     indata = args.indata
-    data = []
+    output = []
     if verbose:
         print("Job: {}".format(job))
-        if isinstance(indata, file):
-            print("Input:"),
-            for line in indata.readlines():
-                data.append(line)
-                print(data)
-        else:
-            print("Data: {}".format(indata))
-            data = indata
 
-    compiled = job + data
-    # print(compiled)
-    process = subprocess.Popen(compiled)
+    if isinstance(indata, file):
+        for line in indata.readlines():
+            output.append(line)
+    else:
+        output = indata
+
+    compiled = job + output
+    process = subprocess.Popen(compiled, stdout=sys.stdout, stderr=sys.stderr)
+    process.communicate()
+    process.wait()
 
 
 if __name__ == "__main__":