annotate data_manager/pangolearn_dm.py @ 5:29c738066906 draft default tip

"planemo upload for repository https://github.com/galaxyproject/tools-iuc/tree/master/data_managers/data_manager_pangolearn commit b12eb7c784ce8e0f9b775530e710e47088711f76"

1 #!/usr/bin/env python
2 
3 import argparse
4 import datetime
5 import json
6 import operator
7 import os
8 import shutil
9 import sys
10 import tarfile
11 
12 import requests
13 
14 
15 def get_model_list(
16     existing_release_tags,
17     url="https://api.github.com/repos/cov-lineages/pangoLEARN/releases"
18 ):
19     page_num = 0
20     while True:
21         page_num += 1
22         response = requests.get(url + f'?page={page_num}')
23         if response.status_code == 200:
24             release_list_chunk = json.loads(response.text)
25             if not release_list_chunk:
26                 # past the last page of results
27                 return
28             for e in release_list_chunk:
29                 if e["tag_name"] in existing_release_tags:
30                     continue
31                 if e["prerelease"]:
32                     continue
33                 yield dict(
34                     tag_name=e["tag_name"],
35                     name=e["name"],
36                     date=parse_date(e["tag_name"]),
37                     tarball_url=e["tarball_url"],
38                 )
39         else:
40             response.raise_for_status()
41 
42 
43 def filter_by_date(existing_release_tags, start_date=None, end_date=None):
44     ret = []
45     for release in get_model_list(existing_release_tags):
46         if start_date and release["date"] < start_date:
47             break
48         if not end_date or release["date"] <= end_date:
49             ret.append(release)
50 
51     return ret
52 
53 
54 def download_and_unpack(url, output_directory, v3datatree=True):
55     response = requests.get(url)
56     if response.status_code == 200:
57         tmp_filename = url.split("/")[-1]
58         tmpfile = open(tmp_filename, "wb")
59         tmpfile.write(response.content)
60         tmpfile.close()
61         shutil.copy(tmp_filename, "/tmp")
62         tf = tarfile.open(tmp_filename)
63         pl_path = os.path.join(output_directory, tf.next().name)
64         tf.extractall(output_directory)
65         os.unlink(tmp_filename)
66         pangolearn_unpacked_dir = os.path.join(pl_path, "pangoLEARN")
67         if v3datatree:
68             # pangolin v3 expects a datadir with the entire pangoLEARN
69             # subfolder in it.
70             # In addition, it will only use the data if the __init__.py file
71             # contained within that subfolder declares a __version__ that is
72             # newer than the version installed with pangolin.
73             pangolearn_dir = os.path.join(
74                 output_directory, tmp_filename, "pangoLEARN"
75             )
76             os.mkdir(os.path.dirname(pangolearn_dir))
77             # rewrite the __init__.py file and make the __version__ string
78             # appear newer than anything that might come with pangolin by
79             # prepending a "v" to it.
80             pangolearn_init = open(
81                 os.path.join(pangolearn_unpacked_dir, "__init__.py")
82             ).readlines()
83             with open(
84                 os.path.join(pangolearn_unpacked_dir, "__init__.py"), "w"
85             ) as o:
86                 for line in pangolearn_init:
87                     if line.startswith('__version__ = "'):
88                         line = line.replace(
89                             '__version__ = "', '__version__ = "v'
90                         )
91                     o.write(line)
92         else:
93             # Earlier versions of pangolin expect a datadir with just the
94             # contents of the downloaded pangoLEARN subfolder in it and don't
95             # care about the declared version in __init__.py.
96             pangolearn_dir = os.path.join(
97                 output_directory, tmp_filename
98             )
99         os.rename(
00             pangolearn_unpacked_dir,
01             pangolearn_dir
02         )
03         shutil.rmtree(pl_path)
04         return tmp_filename
05     else:
06         response.raise_for_status()
07 
08 
09 def parse_date(d):
10     # Tries to parse the first 10 chars of d as a date, which currently
11     # succeeds for all pangolearn model releases.
12     return datetime.datetime.strptime(d[:10], "%Y-%m-%d")
13 
14 
15 if __name__ == "__main__":
16 
17     parser = argparse.ArgumentParser()
18     parser.add_argument("--testmode", default=False, action="store_true")
19     parser.add_argument("--latest", default=False, action="store_true")
20     parser.add_argument("--start_date", type=parse_date)
21     parser.add_argument("--end_date", type=parse_date)
22     parser.add_argument("--overwrite", default=False, action="store_true")
23     parser.add_argument('--pangolearn_format_version')
24     parser.add_argument("datatable_name")
25     parser.add_argument("galaxy_datamanager_filename")
26     args = parser.parse_args()
27 
28     if args.testmode:
29         releases = filter_by_date([], start_date=args.start_date, end_date=args.end_date)
30         for release in releases:
31             print(release["tag_name"], release["tarball_url"].split("/")[-1], release["date"])
32         sys.exit(0)
33 
34     with open(args.galaxy_datamanager_filename) as fh:
35         config = json.load(fh)
36 
37     output_directory = config.get("output_data", [{}])[0].get("extra_files_path", None)
38     data_manager_dict = {}
39     data_manager_dict["data_tables"] = config.get("data_tables", {})
40     data_manager_dict["data_tables"][args.datatable_name] = data_manager_dict[
41         "data_tables"
42     ].get(args.datatable_name, [])
43 
44     # NOTE: the data_manager_dict["data_tables"][args.datatable_name] is not actually populated with the
45     # contents of the existing data table, so the "no-overwrite" logic and the
46     # only-download-what-we-don't-have logic does not in fact work. It is left but unused for now.
47     if not args.overwrite:
48         existing_release_tags = set(
49             [
50                 el["value"]
51                 for el in data_manager_dict["data_tables"][args.datatable_name]
52             ]
53         )
54     else:
55         existing_release_tags = set()
56     if args.latest:
57         releases = [next(get_model_list(existing_release_tags))]
58     else:
59         releases = filter_by_date(
60             existing_release_tags, start_date=args.start_date, end_date=args.end_date
61         )
62     releases_to_download = [
63         release
64         for release in releases
65         if release["tag_name"] not in existing_release_tags
66     ]
67     for release in releases_to_download:
68         if args.pangolearn_format_version is not None:
69             version = args.pangolearn_format_version
70         else:
71             # 2021-05-27 was the first release of pangoLEARN for pangolin 3, which changed DB format
72             if release["date"] >= datetime.datetime(2021, 5, 27):
73                 version = '3.0'
74             else:
75                 version = '1.0'
76         fname = download_and_unpack(
77             release["tarball_url"],
78             output_directory,
79             v3datatree=version == '3.0'
80         )
81         data_manager_dict["data_tables"][args.datatable_name].append(
82             dict(
83                 value=release["tag_name"],
84                 description=release["name"],
85                 format_version=version,
86                 path=output_directory + "/" + fname,
87             )
88         )
89     data_manager_dict["data_tables"][args.datatable_name].sort(
90         key=operator.itemgetter("value"), reverse=True
91     )
92     with open(args.galaxy_datamanager_filename, "w") as fh:
93         json.dump(data_manager_dict, fh, indent=2, sort_keys=True)