docker-stacks/tagging/update_wiki.py

#!/usr/bin/env python3
# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.
import argparse
import datetime
import logging
import shutil
from pathlib import Path

import plumbum
from dateutil import relativedelta

git = plumbum.local["git"]

LOGGER = logging.getLogger(__name__)

YEAR_TABLE_HEADER = """\
## {year}

| Month                  | Builds | Images | Commits |
| ---------------------- | ------ | ------ | ------- |
"""


def build_monthly_table_line(year_month_file: Path) -> str:
    year_month_file_content = year_month_file.read_text()
    images = year_month_file_content.count("Build manifest")
    builds = sum(
        "jupyter/base-notebook" in line and "aarch64" not in line
        for line in year_month_file_content.split("\n")
    )
    year_month = year_month_file.stem
    current_month = datetime.date(
        year=int(year_month[:4]), month=int(year_month[5:]), day=1
    )
    next_month = current_month + relativedelta.relativedelta(months=1)
    future = (
        git["log", "--oneline", "--since", current_month, "--until", next_month]
        & plumbum.BG
    )
    future.wait()
    commits = len(future.stdout.splitlines())

    return f"| [`{year_month}`](./{year_month}) | {builds: <6} | {images: <6} | {commits: <7} |\n"


def regenerate_home_wiki_page(wiki_dir: Path) -> None:
    YEAR_MONTHLY_TABLES = "<!-- YEAR_MONTHLY_TABLES -->\n"

    wiki_home_file = wiki_dir / "Home.md"
    wiki_home_content = wiki_home_file.read_text()

    assert YEAR_MONTHLY_TABLES in wiki_home_content
    wiki_home_content = wiki_home_content[
        : wiki_home_content.find(YEAR_MONTHLY_TABLES) + len(YEAR_MONTHLY_TABLES)
    ]

    for year_dir in sorted((wiki_dir / "monthly-files").glob("*"), reverse=True):
        wiki_home_content += "\n" + YEAR_TABLE_HEADER.format(year=year_dir.name)
        for year_month_file in sorted(year_dir.glob("*.md"), reverse=True):
            wiki_home_content += build_monthly_table_line(year_month_file)

    wiki_home_file.write_text(wiki_home_content)
    LOGGER.info("Updated Home page")


def update_monthly_wiki_page(
    wiki_dir: Path, year_month: str, build_history_line: str
) -> None:
    MONTHLY_PAGE_HEADER = f"""\
# Images built during {year_month}

| Date | Image | Links |
| - | - | - |
"""
    year = year_month[:4]
    monthly_page = wiki_dir / "monthly-files" / year / (year_month + ".md")
    if not monthly_page.exists():
        monthly_page.write_text(MONTHLY_PAGE_HEADER)
        LOGGER.info(f"Created monthly page: {monthly_page.relative_to(wiki_dir)}")

    monthly_page_content = monthly_page.read_text()
    assert MONTHLY_PAGE_HEADER in monthly_page_content
    monthly_page_content = monthly_page_content.replace(
        MONTHLY_PAGE_HEADER, MONTHLY_PAGE_HEADER + build_history_line + "\n"
    )
    monthly_page.write_text(monthly_page_content)
    LOGGER.info(f"Updated monthly page: {monthly_page.relative_to(wiki_dir)}")


def get_manifest_timestamp(manifest_file: Path) -> str:
    file_content = manifest_file.read_text()
    TIMESTAMP_PREFIX = "Build timestamp: "
    TIMESTAMP_LENGTH = 20
    timestamp = file_content[
        file_content.find(TIMESTAMP_PREFIX) + len(TIMESTAMP_PREFIX) :
    ][:TIMESTAMP_LENGTH]
    # Should be good enough till year 2100
    assert timestamp.startswith("20"), timestamp
    assert timestamp.endswith("Z"), timestamp
    return timestamp


def get_manifest_year_month(manifest_file: Path) -> str:
    return get_manifest_timestamp(manifest_file)[:7]


def remove_old_manifests(wiki_dir: Path) -> None:
    MAX_NUMBER_OF_MANIFESTS = 4500

    manifest_files: list[tuple[str, Path]] = []
    for file in (wiki_dir / "manifests").rglob("*.md"):
        manifest_files.append((get_manifest_timestamp(file), file))

    manifest_files.sort(reverse=True)
    for _, file in manifest_files[MAX_NUMBER_OF_MANIFESTS:]:
        file.unlink()
        LOGGER.info(f"Removed manifest: {file.relative_to(wiki_dir)}")


def update_wiki(wiki_dir: Path, hist_lines_dir: Path, manifests_dir: Path) -> None:
    LOGGER.info("Updating wiki")

    manifest_files = list(manifests_dir.rglob("*.md"))
    assert manifest_files, "expected to have some manifest files"
    for manifest_file in manifest_files:
        year_month = get_manifest_year_month(manifest_file)
        year = year_month[:4]
        copy_to = wiki_dir / "manifests" / year / year_month / manifest_file.name
        copy_to.parent.mkdir(exist_ok=True)
        shutil.copy(manifest_file, copy_to)
        LOGGER.info(f"Added manifest file: {copy_to.relative_to(wiki_dir)}")

    build_history_line_files = sorted(hist_lines_dir.rglob("*.txt"))
    assert build_history_line_files, "expected to have some build history line files"
    for build_history_line_file in build_history_line_files:
        build_history_line = build_history_line_file.read_text()
        assert build_history_line.startswith("| `")
        year_month = build_history_line[3:10]
        update_monthly_wiki_page(wiki_dir, year_month, build_history_line)

    regenerate_home_wiki_page(wiki_dir)
    remove_old_manifests(wiki_dir)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "--wiki-dir",
        required=True,
        type=Path,
        help="Directory of the wiki repo",
    )
    arg_parser.add_argument(
        "--hist-lines-dir",
        required=True,
        type=Path,
        help="Directory with history lines",
    )
    arg_parser.add_argument(
        "--manifests-dir",
        required=True,
        type=Path,
        help="Directory with manifest files",
    )
    args = arg_parser.parse_args()

    update_wiki(args.wiki_dir, args.hist_lines_dir, args.manifests_dir)