nedrex/manage.py at main · repotrial/nedrex · GitHub

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
#!/usr/bin/env python

from pathlib import Path

import click
import loguru
import pymongo
import semantic_version

client = pymongo.MongoClient("localhost:27020")
db = client["test"]

results_dir = Path("experimental/fastapi/refactored_api/results")
logger = loguru.logger


@click.group()
def cli():
    pass


@click.group(help="Operations related to graphs build by Repotrial DB")
def graphs():
    pass


@click.group(help="Operations related to results stored in Repotrial DB")
def results():
    pass


@results.command(
    help="Removes the result with the specified UID from the target results"
)
@click.option(
    "-c",
    "--collection",
    type=click.Choice(
        [
            "diamond",
            "bicon",
            "closeness",
            "graphs",
            "must",
            "trustrank",
            "graphs",
        ]
    ),
)
@click.option("-u", "--uid", type=str)
def remove(collection, uid):
    coll = db[f"{collection}_"]
    res = coll.delete_one({"uid": uid})
    if not res.deleted_count == 1:
        logger.warning(
            f"uid {uid} does not appear to be in collection {collection}"
        )

    relevant_directory = results_dir / f"{collection}_"
    for f in relevant_directory.iterdir():
        if f.name.startswith(uid):
            logger.info(f"Removing {f}")
            f.unlink()


@graphs.command(help="Removes graphs built before a specified version")
def clean():
    value = click.prompt(
        "Please enter the version you wish to delete up to (not including)",
        type=str,
    )
    confirm_value = click.prompt(
        "Please confirm the version you wish to delete up to (not including)",
        type=str,
    )

    if value != confirm_value:
        print("Mismatch in values -- please check and retry.")
        return

    try:
        v = semantic_version.Version(value)
    except ValueError:
        print("Version given is not valid -- please check an retry.")
        return

    # Get versions that are lower than the specified version.
    coll = db["graphs_"]
    uids = {
        entry["uid"]
        for entry in coll.find()
        if semantic_version.Version(entry["version"]) < v
    }
    graph_dir = results_dir / "graphs_"

    warning_count = 0

    for uid in uids:
        q = graph_dir / f"{uid}.graphml"
        coll.delete_one({"uid": uid})
        if q.exists():
            q.unlink()
            logger.debug(f"{uid} successfully removed")
        else:
            logger.warning(f"{uid} did not have an associated graphml file.")
            warning_count += 1

    logger.info(f"Removed {len(uids)} entries ({warning_count} warnings).")


cli.add_command(graphs)
cli.add_command(results)


if __name__ == "__main__":
    cli()