X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/783aff1187eba225008cd7db6fdd8307e46bc59b..d56b21609ddd2e2e096f5c30e991d24aa213f7f4:/crunch_scripts/split-fastq.py

diff --git a/crunch_scripts/split-fastq.py b/crunch_scripts/split-fastq.py
index f42b2b0ab9..61c384fbf0 100755
--- a/crunch_scripts/split-fastq.py
+++ b/crunch_scripts/split-fastq.py
@@ -1,107 +1,70 @@
 #!/usr/bin/python
+# Copyright (C) The Arvados Authors. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
 
 import arvados
 import re
 import hashlib
+import string
 
-#api = arvados.api('v1')
+api = arvados.api('v1')
 
 piece = 0
 manifest_text = ""
 
-def readline(reader, start):
-    line = ""
+# Look for paired reads
+
+inp = arvados.CollectionReader(arvados.getjobparam('reads'))
+
+manifest_list = []
+
+def nextline(reader, start):
     n = -1
-    while n == -1:
-        r = reader.readfrom(start, 1024)
+    while True:
+        r = reader.readfrom(start, 128)
         if r == '':
             break
         n = string.find(r, "\n")
-        line += r[0:n]
-        start += len(r)
-    return line
-
-def splitfastq(p):
-    for i in xrange(0, len(p)):
-        p[i]["start"] = 0
-        p[i]["end"] = 0
-
-    while True:
-        recordsize = [0, 0]
-
-        # read 4 lines starting at "start"
-        for ln in xrange(0, 4):
-            for i in xrange(0, len(p)):
-                r = readline(p[i]["reader"], p[i]["start"])
-                if r == '':
-                    return
-                recordsize[i] += len(r)
-
-        splitnow = False
-        for i in xrange(0, len(p)):
-            if ((p[i]["end"] - p[i]["start"]) + recordsize[i]) >= arvados.BLOCKSIZE:
-                splitnow = True
-
-        if splitnow:
-            for i in xrange(0, len(p)):
-                manifest = []
-                manifest.extend(str(piece))
-                manifest.extend([d[LOCATOR] for d in p["reader"]._stream._data_locators])
-                manifest.extend(["{}:{}:{}".format(seg[LOCATOR], seg[BLOCKSIZE], self.name().replace(' ', '\\040')) for seg in arvados.locators_and_ranges(p[i]["reader"].segments, p[i]["start"], p[i]["end"] - p[i]["start"])])
-                global manifest_text
-                manifest_text += manifest.join(" ") + "\n"
-                p[i]["start"] = p[i]["end"]
+        if n > -1:
+            break
         else:
-            for i in xrange(0, len(p)):
-                p[i]["end"] += recordsize[i]
-
-def put_manifest(manifest_text, sources=[]):
-    crm = arvados.CollectionReader(manifest_text)
+            start += 128
+    return n
 
-    combined = crm.manifest_text(strip=True)
-
-    m = hashlib.new('md5')
-    m.update(combined)
-
-    uuid = "{}+{}".format(m.hexdigest(), len(combined))
-
-    collection = arvados.api().collections().create(
-        body={
-            'uuid': uuid,
-            'manifest_text': crm.manifest_text(),
-        }).execute()
-
-    for s in sources:
-        l = arvados.api().links().create(body={
-            "link": {
-                "tail_uuid": s,
-                "head_uuid": uuid,
-                "link_class": "provenance",
-                "name": "provided"
-            }}).execute()
-
-    return uuid
-
-
-# Look for pairs
-
-#inp = arvados.CollectionReader(arvados.get_job_param('input'))
-
-with open("/home/peter/manifest") as f:
-    inp = arvados.CollectionReader(f.read())
-
-prog = re.compile("(.*?)_1.fastq$")
+prog = re.compile(r'(.*?)(_[12])?\.fastq(\.gz)?$')
 
+# Look for fastq files
 for s in inp.all_streams():
-    if s.name() == ".":
-        for f in s.all_files():
-            result = prog.match(f.name())
-            print f.name()
-            if result != None:
-                p = [{}, {}]
-                p[0]["reader"] = s.file(f)
-                p[1]["reader"] = s.file(prog.group(1) + "_2.fastq")
-
-print manifest_text
-
-#arvados.current_task().set_output(put_manifest(manifest_text))
+    for f in s.all_files():
+        name_pieces = prog.match(f.name())
+        if name_pieces is not None:
+            if s.name() != ".":
+                # The downstream tool (run-command) only iterates over the top
+                # level of directories so if there are fastq files in
+                # directories in the input, the choice is either to forget
+                # there are directories (which might lead to name conflicts) or
+                # just fail.
+                print >>sys.stderr, "fastq must be at the root of the collection"
+                sys.exit(1)
+
+            p = None
+            if name_pieces.group(2) is not None:
+                if name_pieces.group(2) == "_1":
+                    p = [{}, {}]
+                    p[0]["reader"] = s.files()[name_pieces.group(0)]
+                    p[1]["reader"] = s.files()[name_pieces.group(1) + "_2.fastq" + (name_pieces.group(3) if name_pieces.group(3) else '')]
+            else:
+                p = [{}]
+                p[0]["reader"] = s.files()[name_pieces.group(0)]
+
+            if p is not None:
+                for i in xrange(0, len(p)):
+                    m = p[i]["reader"].as_manifest().split()
+                    m[0] = "./_" + str(piece)
+                    manifest_list.append(m)
+                piece += 1
+
+manifest_text = "\n".join(" ".join(m) for m in manifest_list) + "\n"
+
+arvados.current_task().set_output(manifest_text)