pytorch

1
#!/usr/bin/env python3
2

3
import sys
4
from pathlib import Path
5
from typing import Any, cast, Dict, List, Set
6

7
import yaml
8

9

10
GITHUB_DIR = Path(__file__).parent.parent
11

12

13
def get_workflows_push_tags() -> Set[str]:
14
    "Extract all known push tags from workflows"
15
    rc: Set[str] = set()
16
    for fname in (GITHUB_DIR / "workflows").glob("*.yml"):
17
        with fname.open("r") as f:
18
            wf_yml = yaml.safe_load(f)
19
        # "on" is alias to True in yaml
20
        on_tag = wf_yml.get(True, None)
21
        push_tag = on_tag.get("push", None) if isinstance(on_tag, dict) else None
22
        tags_tag = push_tag.get("tags", None) if isinstance(push_tag, dict) else None
23
        if isinstance(tags_tag, list):
24
            rc.update(tags_tag)
25
    return rc
26

27

28
def filter_ciflow_tags(tags: Set[str]) -> List[str]:
29
    "Return sorted list of ciflow tags"
30
    return sorted(
31
        tag[:-2] for tag in tags if tag.startswith("ciflow/") and tag.endswith("/*")
32
    )
33

34

35
def read_probot_config() -> Dict[str, Any]:
36
    with (GITHUB_DIR / "pytorch-probot.yml").open("r") as f:
37
        return cast(Dict[str, Any], yaml.safe_load(f))
38

39

40
def update_probot_config(labels: Set[str]) -> None:
41
    orig = read_probot_config()
42
    orig["ciflow_push_tags"] = filter_ciflow_tags(labels)
43
    with (GITHUB_DIR / "pytorch-probot.yml").open("w") as f:
44
        yaml.dump(orig, f, indent=4, sort_keys=False)
45

46

47
if __name__ == "__main__":
48
    from argparse import ArgumentParser
49

50
    parser = ArgumentParser("Validate or update list of tags")
51
    parser.add_argument("--validate-tags", action="store_true")
52
    args = parser.parse_args()
53
    pushtags = get_workflows_push_tags()
54
    if args.validate_tags:
55
        config = read_probot_config()
56
        ciflow_tags = set(filter_ciflow_tags(pushtags))
57
        config_tags = set(config["ciflow_push_tags"])
58
        if config_tags != ciflow_tags:
59
            print("Tags mismatch!")
60
            if ciflow_tags.difference(config_tags):
61
                print(
62
                    "Reference in workflows but not in config",
63
                    ciflow_tags.difference(config_tags),
64
                )
65
            if config_tags.difference(ciflow_tags):
66
                print(
67
                    "Reference in config, but not in workflows",
68
                    config_tags.difference(ciflow_tags),
69
                )
70
            print(f"Please run {__file__} to remediate the difference")
71
            sys.exit(-1)
72
        print("All tags are listed in pytorch-probot.yml")
73
    else:
74
        update_probot_config(pushtags)
75