annotate data_manager/data_manager_plant_tribes_scaffolds_download.py @ 1:80b0bd65cbfb draft

Uploaded
author iuc
date Wed, 29 Mar 2017 12:08:57 -0400
parents 4c96b684f0fd
children 5833ef61c1f8
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
0
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
1 #!/usr/bin/env python
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
2 #
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
3 # Data manager for downloading Plant Tribes scaffolds data.
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
4 import argparse
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
5 import json
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
6 import os
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
7 import shutil
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
8 import sys
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
9 import tarfile
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
10 import urllib2
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
11 import zipfile
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
12
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
13
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
14 DEFAULT_DATA_TABLE_NAMES = ["plant_tribes_scaffolds"]
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
15
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
16
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
17 def add_data_table_entry(data_manager_dict, data_table_name, data_table_entry):
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
18 data_manager_dict['data_tables'] = data_manager_dict.get('data_tables', {})
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
19 data_manager_dict['data_tables'][data_table_name] = data_manager_dict['data_tables'].get(data_table_name, [])
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
20 data_manager_dict['data_tables'][data_table_name].append(data_table_entry)
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
21 return data_manager_dict
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
22
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
23
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
24 def make_directory(dir):
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
25 if not os.path.exists(dir):
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
26 os.makedirs(dir)
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
27
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
28
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
29 def remove_directory(dir):
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
30 if os.path.exists(dir):
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
31 shutil.rmtree(dir)
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
32
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
33
1
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
34 def extract_archive(file_path, work_directory):
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
35 if tarfile.is_tarfile(file_path):
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
36 fh = tarfile.open(file_path, 'r:*')
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
37 elif zipfile.is_zipfile(file_path):
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
38 fh = zipfile.ZipFile(file_path, 'r')
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
39 else:
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
40 return
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
41 fh.extractall(work_directory)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
42
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
43
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
44 def move_files(source_directory, target_directory):
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
45 # Move the files into defined output directory.
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
46 for filename in os.listdir(source_directory):
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
47 shutil.move(os.path.join(source_directory, filename), target_directory)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
48
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
49
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
50 def url_download(url, work_directory):
0
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
51 file_path = os.path.join(work_directory, os.path.basename(url))
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
52 src = None
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
53 dst = None
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
54 try:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
55 req = urllib2.Request(url)
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
56 src = urllib2.urlopen(req)
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
57 dst = open(file_path, 'wb')
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
58 while True:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
59 chunk = src.read(2**10)
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
60 if chunk:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
61 dst.write(chunk)
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
62 else:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
63 break
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
64 except Exception, e:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
65 print >>sys.stderr, str(e)
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
66 finally:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
67 if src:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
68 src.close()
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
69 if dst:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
70 dst.close()
1
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
71 return file_path
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
72
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
73
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
74 def download(target_directory, web_url, config_web_url, description, data_table_names=DEFAULT_DATA_TABLE_NAMES):
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
75 data_manager_dict = {}
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
76 data_table_entry = {}
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
77 # Download the scaffolds data.
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
78 work_directory = os.path.abspath(os.path.join(os.getcwd(), 'scaffolds'))
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
79 make_directory(work_directory)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
80 file_path = url_download(web_url, work_directory)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
81 extract_archive(file_path, work_directory)
0
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
82 os.remove(file_path)
1
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
83 # Move the scaffolds data files into the defined output directory.
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
84 move_files(work_directory, target_directory)
0
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
85 remove_directory(work_directory)
1
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
86 # Populate the data_manager_dict with the scaffolds data entry.
0
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
87 for file_path in os.listdir(target_directory):
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
88 full_path = os.path.abspath(os.path.join(target_directory, file_path))
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
89 entry_name = "%s" % os.path.basename(file_path)
1
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
90 data_table_entry['value'] = entry_name
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
91 data_table_entry['name'] = entry_name
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
92 data_table_entry['path'] = full_path
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
93 data_table_entry['description'] = description
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
94 # Populate the data_manager_dict.
0
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
95 for data_table_name in data_table_names:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
96 data_manager_dict = add_data_table_entry(data_manager_dict, data_table_name, data_table_entry)
1
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
97 # Download the default configuration files.
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
98 work_directory = os.path.abspath(os.path.join(os.getcwd(), 'configs'))
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
99 make_directory(work_directory)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
100 file_path = url_download(config_web_url, work_directory)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
101 extract_archive(file_path, work_directory)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
102 os.remove(file_path)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
103 # Move the default configuration files into the defined output directory.
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
104 source_configs_directory = os.path.join(work_directory, entry_name)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
105 target_configs_directory = os.path.join(target_directory, entry_name)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
106 move_files(source_configs_directory, target_configs_directory)
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
107 remove_directory(work_directory)
0
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
108 return data_manager_dict
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
109
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
110
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
111 parser = argparse.ArgumentParser()
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
112 parser.add_argument('--description', dest='description', default=None, help='Description')
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
113 parser.add_argument('--name', dest='name', help='Data table entry unique ID')
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
114 parser.add_argument('--out_file', dest='out_file', help='JSON output file')
1
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
115 parser.add_argument('--web_url', dest='web_url', help='URL for downloading scaffolds')
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
116 parser.add_argument('--config_web_url', dest='config_web_url', help='URL for downloading default configs')
0
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
117
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
118 args = parser.parse_args()
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
119
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
120 # Some magic happens with tools of type "manage_data" in that the output
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
121 # file contains some JSON data that allows us to define the target directory.
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
122 params = json.loads(open(args.out_file).read())
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
123 target_directory = params['output_data'][0]['extra_files_path']
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
124 make_directory(target_directory)
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
125
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
126 if args.description is None:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
127 description = ''
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
128 else:
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
129 description = args.description.strip()
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
130
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
131 # Get the scaffolds data.
1
80b0bd65cbfb Uploaded
iuc
parents: 0
diff changeset
132 data_manager_dict = download(target_directory, args.web_url, args.config_web_url, description)
0
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
133 # Write the JSON output dataset.
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
134 fh = open(args.out_file, 'wb')
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
135 fh.write(json.dumps(data_manager_dict))
4c96b684f0fd Uploaded
iuc
parents:
diff changeset
136 fh.close()