|
0
|
1 import argparse
|
|
|
2 import logging
|
|
|
3 import os
|
|
|
4 import re
|
|
|
5 import shutil
|
|
|
6 import sys
|
|
|
7 import tarfile
|
|
|
8 import tempfile
|
|
|
9 import zipfile
|
|
|
10
|
|
|
11 import magic
|
|
|
12
|
|
|
13 imgt_file_regex = re.compile("^\d+_[^P]")
|
|
|
14
|
|
|
15
|
|
|
16 def sniff_imgt_type(input_file):
|
|
|
17 m = magic.Magic()
|
|
|
18 file_type = m.from_file(input_file)
|
|
|
19 logging.debug("File type of {0} is {1}".format(input_file, file_type))
|
|
|
20 return file_type.split(" ")[0]
|
|
|
21
|
|
|
22
|
|
|
23 def unpack_imgt_zip(input_file, output_dir):
|
|
|
24 imgt_type = sniff_imgt_type(input_file)
|
|
|
25 if imgt_type == "Zip":
|
|
|
26 with zipfile.ZipFile(input_file) as inf:
|
|
|
27 inf.extractall(output_dir)
|
|
|
28 elif imgt_type == "XZ":
|
|
|
29 with tarfile.open(input_file) as inf:
|
|
|
30 inf.extractall(output_dir)
|
|
|
31 else:
|
|
|
32 raise IOError("Unsuppported file type: {0}".format(imgt_type))
|
|
|
33 logging.debug("Extracted {0} to {1}".format(input_file, output_dir))
|
|
|
34 check = os.listdir(output_dir)
|
|
|
35 if len(check) == 1:
|
|
|
36 check = os.path.join(output_dir, check[0])
|
|
|
37 if os.path.isdir(check):
|
|
|
38 logging.info("{0} is an older IMGT zip, removing extra dir".format(input_file))
|
|
|
39 files = os.listdir(check)
|
|
|
40 for file in files:
|
|
|
41 new_file = os.path.join(output_dir, file)
|
|
|
42 file = os.path.join(check, file)
|
|
|
43 shutil.move(file, new_file)
|
|
|
44 shutil.rmtree(check)
|
|
|
45
|
|
|
46
|
|
1
|
47 def filter_imgt_file(old_file, new_file, column, fltr):
|
|
|
48 logging.debug("Filtering {0} with {1}".format(old_file, fltr))
|
|
0
|
49 first = True
|
|
|
50 total = 0
|
|
|
51 remain = 0
|
|
|
52 with open(old_file, 'r') as of, open(new_file, 'w') as nf:
|
|
|
53 column_index = -1
|
|
|
54 for line in of:
|
|
|
55 splt = line.split("\t")
|
|
|
56 if first:
|
|
|
57 column_index = splt.index(column)
|
|
|
58 first = False
|
|
|
59 nf.write(line)
|
|
|
60 continue
|
|
|
61 total += 1
|
|
1
|
62 if len(splt) > column_index and splt[column_index].find(fltr) != -1:
|
|
0
|
63 remain += 1
|
|
|
64 nf.write(line)
|
|
|
65 return total, remain
|
|
|
66
|
|
|
67
|
|
|
68 def all_same_in_list(l):
|
|
|
69 return all(l[0] == x for x in l[1:])
|
|
|
70
|
|
|
71
|
|
1
|
72 def filter_imgt_dir(imgt_dir, locus):
|
|
|
73 logging.info("Working on {0}".format(locus))
|
|
0
|
74 imgt_files = [f for f in os.listdir(imgt_dir) if imgt_file_regex.match(f)]
|
|
|
75 tmp_file = os.path.join(imgt_dir, "tmp.txt")
|
|
|
76 totals = []
|
|
|
77 remains = []
|
|
|
78 for imgt_file in imgt_files:
|
|
|
79 imgt_file = os.path.join(imgt_dir, imgt_file)
|
|
1
|
80 total, remain = filter_imgt_file(imgt_file, tmp_file, "V-GENE and allele", locus)
|
|
0
|
81 totals.append(total)
|
|
|
82 remains.append(remain)
|
|
|
83 logging.debug("{0} rows, {1} after filtering".format(total, remain))
|
|
|
84 shutil.move(tmp_file, imgt_file)
|
|
|
85 if not (all_same_in_list(totals) and all_same_in_list(remains)):
|
|
1
|
86 logging.warning("Not all files had the same number of sequences remaining for {0}".format(imgt_dir))
|
|
0
|
87 return totals[0], remains[0]
|
|
|
88
|
|
|
89
|
|
|
90 def make_new_xz_file(input_dir, output_file):
|
|
|
91 logging.info("Creating new IMGT zip for {0} at {1}".format(input_dir, output_file))
|
|
|
92 imgt_files = [f for f in os.listdir(input_dir)]
|
|
|
93 with tarfile.open(output_file, 'w:xz') as out:
|
|
|
94 for imgt_file in imgt_files:
|
|
|
95 logging.debug("Writing {0} to new IMGT zip".format(imgt_file))
|
|
|
96 imgt_file = os.path.join(input_dir, imgt_file)
|
|
|
97 out.add(imgt_file, arcname=os.path.basename(imgt_file))
|
|
|
98
|
|
|
99
|
|
|
100 def main():
|
|
|
101 parser = argparse.ArgumentParser()
|
|
1
|
102 parser.add_argument("--input", help="The input IMGT file", required=True)
|
|
|
103 parser.add_argument("--output-ig", help="The output file for new IMGT ZIP with just IG sequences", default="None")
|
|
|
104 parser.add_argument("--output-igh", help="The output file for new IMGT ZIP with just IGH sequences", default="None")
|
|
|
105 parser.add_argument("--output-igk", help="The output file for new IMGT ZIP with just IGK sequences", default="None")
|
|
|
106 parser.add_argument("--output-igl", help="The output file for new IMGT ZIP with just IGL sequences", default="None")
|
|
|
107 parser.add_argument("--output-tr", help="The output file for new IMGT ZIP with just TR sequences", default="None")
|
|
|
108 parser.add_argument("--output-tra", help="The output file for new IMGT ZIP with just TRA sequences", default="None")
|
|
|
109 parser.add_argument("--output-trb", help="The output file for new IMGT ZIP with just TRB sequences", default="None")
|
|
|
110 parser.add_argument("--output-trd", help="The output file for new IMGT ZIP with just TRD sequences", default="None")
|
|
|
111 parser.add_argument("--output-trg", help="The output file for new IMGT ZIP with just TRG sequences", default="None")
|
|
0
|
112
|
|
|
113 logging.basicConfig(filename="./log.html", level=logging.DEBUG, format="%(asctime)s: %(message)s <br />",
|
|
|
114 datefmt='%Y/%m/%d %H:%M:%S')
|
|
|
115 logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
|
|
|
116 logging.info("Started IMGT locus split")
|
|
|
117
|
|
|
118 args = parser.parse_args()
|
|
|
119 input_file = args.input
|
|
1
|
120 output_ig = args.output_ig
|
|
|
121 output_igh = args.output_igh
|
|
|
122 output_igk = args.output_igk
|
|
|
123 output_igl = args.output_igl
|
|
|
124 output_tr = args.output_tr
|
|
|
125 output_tra = args.output_tra
|
|
|
126 output_trb = args.output_trb
|
|
|
127 output_trd = args.output_trd
|
|
|
128 output_trg = args.output_trg
|
|
|
129
|
|
|
130 loci = {
|
|
|
131 "IG": output_ig,
|
|
|
132 "IGH": output_igh,
|
|
|
133 "IGK": output_igk,
|
|
|
134 "IGL": output_igl,
|
|
|
135 "TR": output_tr,
|
|
|
136 "TRA": output_tra,
|
|
|
137 "TRB": output_trb,
|
|
|
138 "TRD": output_trd,
|
|
|
139 "TRG": output_trg
|
|
|
140 }
|
|
|
141
|
|
|
142 loci_to_filter = {}
|
|
0
|
143
|
|
|
144 logging.debug("All Parameters:")
|
|
|
145 logging.debug("Input: {0}".format(input_file))
|
|
1
|
146 for locus, path in loci.items():
|
|
|
147 logging.debug("{0}: {1}".format(locus, path))
|
|
|
148 if path != "None" and os.path.isdir(os.path.split(path)[0]):
|
|
|
149 loci_to_filter[locus] = path
|
|
0
|
150
|
|
1
|
151 if len(loci_to_filter) == 0:
|
|
0
|
152 raise Exception("No locus selected, nothing to do")
|
|
|
153
|
|
1
|
154 logging.info("Parameters:")
|
|
|
155 for locus, path in loci_to_filter.items():
|
|
|
156 logging.info("{0}: {1}".format(locus, path))
|
|
|
157
|
|
0
|
158 work_dir = tempfile.mkdtemp()
|
|
|
159 original_files_dir = os.path.join(work_dir, "original")
|
|
|
160 os.mkdir(original_files_dir)
|
|
|
161
|
|
|
162 unpack_imgt_zip(input_file, original_files_dir)
|
|
|
163
|
|
1
|
164 for locus, path in loci_to_filter.items():
|
|
|
165 locus_dir = os.path.join(work_dir, locus)
|
|
|
166 shutil.copytree(original_files_dir, locus_dir)
|
|
|
167 total, remain = filter_imgt_dir(locus_dir, locus)
|
|
|
168 logging.info("{0}\t{1}\t{2}\t{3}".format(locus, path, total, remain))
|
|
0
|
169
|
|
1
|
170 make_new_xz_file(locus_dir, loci_to_filter[locus])
|
|
0
|
171
|
|
|
172
|
|
|
173 if __name__ == "__main__":
|
|
|
174 main()
|
