.ci/tagger.py

import json
import glob
import mmap
import sys


def get_notebooks(path: str):
    return glob.glob(f"{path}/*/[0-9]*.ipynb")


def get_tags(path: str):
    return json.load(open(path))


def find_tags_for_notebook(notebook_path: str, tags: dict):
    nb_tags = []
    with open(notebook_path) as file:
        f = mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ)
        for tag, keywords in tags.items():
            for keyword in keywords:
                if f.find(bytes(keyword, "utf-8")) != -1:
                    nb_tags.append(tag)
                    break
    return nb_tags


def find_tags_for_all_notebooks(notebooks: list, tags: dict):
    notebooks_tags = {}
    for notebook in notebooks:
        nb_tags = sorted(find_tags_for_notebook(notebook, tags))
        if nb_tags:
            notebooks_tags[notebook.split("/")[-1].split(".")[0]] = nb_tags
    return notebooks_tags


if __name__ == "__main__":
    if len(sys.argv) == 1:
        notebooks_paths = sorted(get_notebooks("notebooks"))
        tags = get_tags(".ci/keywords.json")["tags"]
    else:
        notebooks_paths = sorted(get_notebooks("/".join(sys.argv[1].split("/")[:-2])))
        tags = get_tags(sys.argv[2])["tags"]
    all_notebooks_tags = find_tags_for_all_notebooks(notebooks_paths, tags)
    print(json.dumps(all_notebooks_tags, indent=4))