11167: Refactored tests to use new helper function.
[arvados.git] / crunch_scripts / split-fastq.py
1 #!/usr/bin/python
2 # Copyright (C) The Arvados Authors. All rights reserved.
3 #
4 # SPDX-License-Identifier: Apache-2.0
5
6 import arvados
7 import re
8 import hashlib
9 import string
10
11 api = arvados.api('v1')
12
13 piece = 0
14 manifest_text = ""
15
16 # Look for paired reads
17
18 inp = arvados.CollectionReader(arvados.getjobparam('reads'))
19
20 manifest_list = []
21
22 def nextline(reader, start):
23     n = -1
24     while True:
25         r = reader.readfrom(start, 128)
26         if r == '':
27             break
28         n = string.find(r, "\n")
29         if n > -1:
30             break
31         else:
32             start += 128
33     return n
34
35 prog = re.compile(r'(.*?)(_[12])?\.fastq(\.gz)?$')
36
37 # Look for fastq files
38 for s in inp.all_streams():
39     for f in s.all_files():
40         name_pieces = prog.match(f.name())
41         if name_pieces is not None:
42             if s.name() != ".":
43                 # The downstream tool (run-command) only iterates over the top
44                 # level of directories so if there are fastq files in
45                 # directories in the input, the choice is either to forget
46                 # there are directories (which might lead to name conflicts) or
47                 # just fail.
48                 print >>sys.stderr, "fastq must be at the root of the collection"
49                 sys.exit(1)
50
51             p = None
52             if name_pieces.group(2) is not None:
53                 if name_pieces.group(2) == "_1":
54                     p = [{}, {}]
55                     p[0]["reader"] = s.files()[name_pieces.group(0)]
56                     p[1]["reader"] = s.files()[name_pieces.group(1) + "_2.fastq" + (name_pieces.group(3) if name_pieces.group(3) else '')]
57             else:
58                 p = [{}]
59                 p[0]["reader"] = s.files()[name_pieces.group(0)]
60
61             if p is not None:
62                 for i in xrange(0, len(p)):
63                     m = p[i]["reader"].as_manifest().split()
64                     m[0] = "./_" + str(piece)
65                     manifest_list.append(m)
66                 piece += 1
67
68 manifest_text = "\n".join(" ".join(m) for m in manifest_list) + "\n"
69
70 arvados.current_task().set_output(manifest_text)