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 logging
import shutil
from pathlib import Path

LOGGER = logging.getLogger(__name__)


def update_home_wiki_page(wiki_dir: Path, year_month: str) -> None:
    YEAR_MONTHLY_TABLES = "<!-- YEAR_MONTHLY_TABLES -->\n"

    TABLE_HEADER = """\
| Month                  |
| ---------------------- |
"""

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

    year = year_month[:4]
    year_header = f"## {year}\n"
    if year_header not in wiki_home_content:
        assert YEAR_MONTHLY_TABLES in wiki_home_content
        wiki_home_content = wiki_home_content.replace(
            YEAR_MONTHLY_TABLES,
            YEAR_MONTHLY_TABLES + f"\n{year_header}\n{TABLE_HEADER}",
        )
        LOGGER.info(f"Updated wiki home page with year header for year: {year}")

    year_month_line = f"| [`{year_month}`](./{year_month}) |\n"
    if year_month_line not in wiki_home_content:
        assert TABLE_HEADER in wiki_home_content
        wiki_home_content = wiki_home_content.replace(
            TABLE_HEADER, TABLE_HEADER + year_month_line
        )
        LOGGER.info(f"Updated wiki home page with month: {year_month}")

    wiki_home_file.write_text(wiki_home_content)


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_home_wiki_page(wiki_dir, year_month)
        update_monthly_wiki_page(wiki_dir, year_month, build_history_line)

    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)