aboutsummaryrefslogtreecommitdiff
path: root/scripts/sampledata.py
blob: 543ec70abfb30f9aef3c4d8024d69c5fad9eca97 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import json
import os
import sys
import time

from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from functools import partial

# Required Evils!
from flask import g
from wqflask import app

from wqflask.api.gen_menu import gen_dropdown_json
from wqflask.show_trait import show_trait
from wqflask.database import database_connection
from wqflask.search_results import SearchResultPage


class UserSessionSimulator():
    def __init__(self, user_id):
        self._user_id = user_id

    @property
    def user_id(self):
        return self._user_id


def dump_sample_data(dataset_name, trait_id):
    """Given a DATASET_NAME e.g. 'BXDPublish' and a TRAIT_ID
    e.g. '10007', dump the sample data as json object"""
    with database_connection() as conn, conn.cursor() as cursor:
        sample_data = {"headers": ["Name", "Value", "SE"], "data": []}

        with app.app_context():
            g.user_session = UserSessionSimulator(None)
            data = show_trait.ShowTrait(
                cursor, user_id=None,
                kw={
                    "trait_id": trait_id,
                    "dataset": dataset_name
                }
            )
            attributes = data.js_data.get("attributes")
            for id_ in attributes:
                sample_data["headers"].append(attributes[id_].name)
            for sample in data.js_data.get("sample_lists")[0]:
                sample_data["data"].append(
                    [
                        sample.name,
                        sample.value or 'x',
                        sample.variance or 'x',
                        *[str(sample.extra_attributes.get(str(key), "x"))
                          for key in attributes],
                    ])
            return sample_data


def fetch_all_traits(species, group, type_, dataset):
    with app.app_context():
        g.user_session = UserSessionSimulator(None)
        for result in SearchResultPage({
                "species": species,
                "group": group,
                "type": type_,
                "dataset": dataset,
                "search_terms_or": "*",
        }).trait_list:
            yield result.get('name') or result.get('display_name')


def dump_json(base_dir, dataset_name, trait):
    print(f"\033[FDumping: {dataset_name}/{trait}]")
    with open(os.path.join(base_dir, f"{trait}.json"), "w") as f:
        json.dump(dump_sample_data(dataset_name, trait), f)


def dump_dataset(target_dir, species, group_name, dataset_type, dataset):
    start = time.perf_counter()
    dataset_name = dataset[1]
    if not os.path.isdir(
            BASE_DIR := os.path.join(
                target_dir,
                dataset_name
            )
    ):
        os.makedirs(BASE_DIR)
    _l = len(f"Dumping {dataset_name} into {target_dir}:")
    print(f"""
{'='*_l}
Dumping {dataset_name} into {sys.argv[1]}:
{'='*_l}
""")
    with ThreadPoolExecutor() as executor:
        executor.map(
            partial(
                dump_json,
                BASE_DIR,
                dataset_name
            ),
            fetch_all_traits(
                species=species,
                group=group_name,
                type_=dataset_type,
                dataset=dataset_name
            )
        )
    print(f"\033[FDONE DUMPING: {BASE_DIR} !!")
    finish = time.perf_counter()
    print(f"It took {finish-start: .2f} second(s) to finish")


def main():
    # Dump all sampledata into a given directory
    with database_connection() as conn:
        for species, group in gen_dropdown_json(conn).get("datasets").items():
            for group_name, type_ in group.items():
                for dataset_type, datasets in type_.items():
                    with ProcessPoolExecutor() as p_exec:
                        p_exec.map(
                            partial(
                                dump_dataset,
                                sys.argv[1],
                                species,
                                group_name,
                                dataset_type
                            ),
                            datasets
                        )


if __name__ == "__main__":
    main()