Mercurial > repos > iuc > data_manager_fetch_refseq

comparison data_manager/fetch_refseq.py @ 0:8b91891ae805 draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload for repository https://github.com/galaxyproject/tools-iuc/tree/master/data_managers/data_manager_fetch_refseq commit a572f1f01161527ff6ed4af05bb2e073a8ca903b
author iuc
date Mon, 01 Oct 2018 15:36:01 -0400
parents
children d58cad5baa70
comparison
equal deleted inserted replaced
-1:000000000000 0:8b91891ae805
1 #!/usr/bin/env python
2
3 from __future__ import division, print_function
4
5 import argparse
6 import functools
7 import gzip
8 import json
9 import os
10 import os.path
11 import sys
12 from datetime import date
13 from multiprocessing import Process, Queue
14
15 import requests
16
17 try:
18 from io import StringIO
19 except ImportError:
20 from StringIO import StringIO
21 # Refseq structure
22 # - Release number
23 # - Divisions
24 # 1. archea
25 # 2. bacteria
26 # 3. fungi
27 # 4. invertebrate
28 # 5. mitochondrion
29 # 6. other
30 # 7. plant
31 # 8. plasmid
32 # 9. plastid
33 # 10. protozoa
34 # 11. vertebrate mammalian
35 # 12. vertebrate other
36 # 13. viral
37 # within each division
38 # DIVNAME.\d+(.\d+)?.(genomic|protein|rna).(fna|gbff|faa|gpff).gz
39 # where fna and faa are FASTA, gbff and gpff are Genbank
40
41
42 def _add_data_table_entry(data_manager_dict, data_table_entry, data_table_name):
43 data_manager_dict['data_tables'] = data_manager_dict.get('data_tables', {})
44 data_manager_dict['data_tables'][data_table_name] = data_manager_dict['data_tables'].get('all_fasta', [])
45 data_manager_dict['data_tables'][data_table_name].append(data_table_entry)
46 return data_manager_dict
47
48
49 def unzip_to(conn, out_dir, output_filename, chunk_size=4096, debug=False, compress=False):
50 input_filename = conn.get()
51 if compress:
52 open_output = gzip.open
53 else:
54 open_output = open
55 with open_output(os.path.join(out_dir, output_filename), 'wb') as output_file:
56 while input_filename != 'STOP':
57 if debug:
58 print('Reading', input_filename, file=sys.stderr)
59 with gzip.open(input_filename, 'rb') as input_file:
60 read_chunk = functools.partial(input_file.read, (chunk_size))
61 for data in iter(read_chunk, b''): # use b'' as a sentinel to stop the loop. note '' != b'' in Python 3
62 output_file.write(data)
63 os.unlink(input_filename)
64 input_filename = conn.get()
65
66
67 def get_refseq_division(division_name, mol_types, output_directory, debug=False, compress=False):
68 base_url = 'https://ftp.ncbi.nlm.nih.gov/refseq/release/'
69 valid_divisions = set(['archea', 'bacteria', 'complete', 'fungi', 'invertebrate', 'mitochondrion', 'other',
70 'plant', 'plasmid', 'plastid', 'protozoa', 'vertebrate_mammalian', 'vertebrate_other', 'viral'])
71 ending_mappings = {
72 'genomic': '.genomic.fna.gz',
73 'protein': '.protein.faa.gz',
74 'rna': 'rna.fna.gz'
75 }
76 assert division_name in valid_divisions, "Unknown division name ({})".format(division_name)
77 for mol_type in mol_types:
78 assert mol_type in ending_mappings, "Unknown molecule type ({})".format(mol_type)
79 if not os.path.exists(output_directory):
80 os.mkdir(output_directory)
81 release_num_file = base_url + 'RELEASE_NUMBER'
82 r = requests.get(release_num_file)
83 release_num = r.text.strip()
84 division_base_url = base_url + division_name
85 if debug:
86 print('Retrieving {}'.format(division_base_url), file=sys.stderr)
87 r = requests.get(division_base_url)
88 listing_text = r.text
89
90 unzip_queues = {}
91 unzip_processes = []
92 final_output_filenames = []
93 for mol_type in mol_types:
94 q = unzip_queues[mol_type] = Queue()
95 output_filename = division_name + '.' + release_num + '.' + mol_type + '.fasta'
96 if compress:
97 output_filename += '.gz'
98 final_output_filenames.append(output_filename)
99 unzip_processes.append(Process(target=unzip_to, args=(q, output_directory, output_filename),
100 kwargs=dict(debug=debug, compress=compress)))
101 unzip_processes[-1].start()
102
103 # sample line: <a href="vertebrate_other.86.genomic.gbff.gz">vertebrate_other.86.genomic.gbff.gz</a> 2018-07-13 00:59 10M
104 for line in StringIO(listing_text):
105 if '.gz' not in line:
106 continue
107 parts = line.split('"')
108 assert len(parts) == 3, "Unexpected line format: {}".format(line.rstrip())
109 filename = parts[1]
110 for mol_type in mol_types:
111 ending = ending_mappings[mol_type]
112 if filename.endswith(ending):
113 if debug:
114 print('Downloading:', filename, ending, mol_type, file=sys.stderr)
115 output_filename = os.path.join(output_directory, filename)
116 with open(output_filename, 'wb') as output_file:
117 r = requests.get(division_base_url + '/' + filename)
118 for chunk in r.iter_content(chunk_size=4096):
119 output_file.write(chunk)
120 conn = unzip_queues[mol_type]
121 conn.put(output_filename)
122
123 for mol_type in mol_types:
124 conn = unzip_queues[mol_type]
125 conn.put('STOP')
126
127 return [release_num, final_output_filenames]
128
129
130 if __name__ == '__main__':
131 parser = argparse.ArgumentParser(description='Download RefSeq databases')
132 parser.add_argument('--debug', default=False, action='store_true', help='Print debugging output to stderr (verbose)')
133 parser.add_argument('--compress', default=False, action='store_true', help='Compress output files')
134 parser.add_argument('--output_directory', default='tmp', help='Directory to write output to')
135 parser.add_argument('--galaxy_datamanager_filename', help='Galaxy JSON format file describing data manager inputs')
136 parser.add_argument('--division_names', help='RefSeq divisions to download')
137 parser.add_argument('--mol_types', help='Molecule types (genomic, rna, protein) to fetch')
138 parser.add_argument('--pin_date', help='Force download date to this version string')
139 args = parser.parse_args()
140
141 division_names = args.division_names.split(',')
142 mol_types = args.mol_types.split(',')
143 if args.galaxy_datamanager_filename is not None:
144 dm_opts = json.loads(open(args.galaxy_datamanager_filename).read())
145 output_directory = dm_opts['output_data'][0]['extra_files_path'] # take the extra_files_path of the first output parameter
146 data_manager_dict = {}
147 else:
148 output_directory = args.output_directory
149 for division_name in division_names:
150 if args.pin_date is not None:
151 today_str = args.pin_date
152 else:
153 today_str = date.today().strftime('%Y-%m-%d') # ISO 8601 date format
154 [release_num, fasta_files] = get_refseq_division(division_name, mol_types, output_directory, args.debug, args.compress)
155 if args.galaxy_datamanager_filename is not None:
156 for i, mol_type in enumerate(mol_types):
157 assert mol_type in fasta_files[i], "Filename does not contain expected mol_type ({}, {})".format(mol_type, fasta_files[i])
158 unique_key = 'refseq_' + division_name + '.' + release_num + '.' + mol_type # note: this is now same as dbkey
159 dbkey = unique_key
160 desc = 'RefSeq ' + division_name + ' Release ' + release_num + ' ' + mol_type + ' (' + today_str + ')'
161 path = os.path.join(output_directory, fasta_files[i])
162 _add_data_table_entry(data_manager_dict=data_manager_dict,
163 data_table_entry=dict(value=unique_key, dbkey=dbkey, name=desc, path=path),
164 data_table_name='all_fasta')
165 open(args.galaxy_datamanager_filename, 'w').write(json.dumps(data_manager_dict, sort_keys=True))