generate a multi-level PyPi index as per PEP 503 (#1)

[PEP 503](https://peps.python.org/pep-0503/) describes a multi-level PyPi index with a root index file with just links to the supported "projects" and project-specific index files with the actual links to release files. Update `generate_index.py` to generate a multi-level index and modify the test accordingly. The legacy links remains for now until we consider a migration path for how to remove them going forward.
pantsbuild · Feb 14, 2025 · d4cd71d · d4cd71d
1 parent 46f3188
commit d4cd71d
Show file tree

Hide file tree

Showing 4 changed files with 145 additions and 29 deletions.
diff --git a/.github/workflows/pages.yaml b/.github/workflows/pages.yaml
@@ -25,14 +25,18 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup Pages
         uses: actions/configure-pages@v3
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.9'
       - name: Generate index
         run: |
-          mkdir -p github-pages/simple
-          pip install pygithub
-          python generate_index.py > github-pages/simple/index.html
+          python -m venv venv
+          venv/bin/pip install -r ./requirements.txt
+          mkdir -p github-pages
+          venv/bin/python generate_index.py --url-path-prefix=/simple github-pages/simple
         env:
           GH_TOKEN: ${{ github.token }}
       - name: Upload artifact

diff --git a/generate_index.py b/generate_index.py
@@ -1,30 +1,136 @@
+from __future__ import annotations
+
+import argparse
+from collections import defaultdict
 import os
+from pathlib import Path
+import sys
+from typing import Any, Iterable
+import github.GitReleaseAsset
+from packaging.utils import parse_wheel_filename
+from packaging.version import Version
+from urllib.parse import urlparse
+from textwrap import dedent
 
 import github
 
+##
+## Output a PEP 503 compliant package repository for Pants wheels.
+## See https://peps.python.org/pep-0503/
+##
+
 
-def main() -> str:
-    gh = github.Github(auth=github.Auth.Token(os.environ["GH_TOKEN"]))
+def get_pants_python_packages(gh: github.Github) -> dict[str, dict[Version, list[Any]]]:
     repo = gh.get_repo("pantsbuild/pants")
-    releases = repo.get_releases()
-    index = "\n".join(
-        [
-            "<html>",
-            "<body>",
-            "<h1>Links for Pantsbuild Wheels</h1>",
-            *(
-                f'<a href="{asset.browser_download_url}">{asset.name}</a><br>'
-                for release in releases
-                if release.tag_name.startswith("release_2")
-                for asset in release.assets
-                if asset.name.endswith(".whl")
-            ),
-            "</body>",
-            "</html>",
-        ]
-    )
-    return index
+    all_releases = repo.get_releases()
+
+    pants_wheel_assets = [
+        asset
+        for release in all_releases
+        if release.tag_name.startswith("release_2")
+        for asset in release.assets
+        if asset.name.endswith(".whl")
+    ]
+
+    packages = defaultdict(lambda: defaultdict(list))
+
+    for asset in pants_wheel_assets:
+        name, version, _build_tag, _tags = parse_wheel_filename(asset.name)
+        packages[name][version].append(asset)
+
+    return packages
+
+
+def _legacy_flat_links(packages: dict[str, dict[Version, list[Any]]]) -> tuple[str, ...]:
+    return [
+        f'<a href="{asset.browser_download_url}">{asset.name}</a><br>'
+        for package_versions in packages.values()
+        for _, version_release_assets in sorted(package_versions.items(), reverse=True)
+        for asset in version_release_assets
+    ]
+
+# http://repository.example.com/simple/PACKAGE_NAME/
+def _write_package_specific_index(output_dir: Path, package_name: str, package_versions: dict[Version, list[Any]]) -> None:
+    package_output_dir = output_dir / package_name
+    package_output_dir.mkdir()
+
+    package_version_keys = sorted(package_versions.keys(), reverse=True)
+
+    with open(package_output_dir / "index.html", "w") as f:
+        f.write(dedent(
+            f"""\
+            <!DOCTYPE html>
+            <html>
+            <body>
+            <h1>Links for Pantsbuild Wheels - {package_name}</h1>
+            <ul>
+            """
+        ))
+
+        for package_version_key in package_version_keys:
+            package_version_assets = package_versions[package_version_key]
+            package_version_assets.sort(key=lambda x: x.name)
+            for asset in package_version_assets:
+                f.write(f"""<li><a href="{asset.browser_download_url}">{asset.name}</a></li>\n""")
+
+        f.write(dedent(
+            """\
+            </ul>
+            </body>
+            </html>
+            """
+        ))
+
+
+
+def main(args):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url-path-prefix", default="/", action="store")
+    parser.add_argument("output_dir", action="store")
+    opts = parser.parse_args(args)
+
+    github_client = github.Github(auth=github.Auth.Token(os.environ["GH_TOKEN"]))
+    packages = get_pants_python_packages(github_client)
+    package_names = sorted(packages.keys())
+
+    prefix = opts.url_path_prefix
+    if prefix and prefix.endswith("/"):
+        prefix = prefix[0:-1]
+
+    output_dir = Path(opts.output_dir)
+    if output_dir.exists():
+        raise Exception(f"Output directory `{output_dir}` already exists.")
+    output_dir.mkdir(parents=True)
+
+    # http://repository.example.com/simple/
+    with open(output_dir / "index.html", "w") as f:
+        f.write(dedent(
+            """\
+            <!DOCTYPE html>
+            <html>
+            <body>
+            <h1>Links for Pantsbuild Wheels</h1>
+            <ul>
+            """
+        ))
+
+        for package_name in package_names:
+            f.write(f"""<li><a href="{prefix}/{package_name}/">{package_name}</a></li>\n""")
+
+        f.write("\n".join(_legacy_flat_links(packages)))
+
+        f.write(dedent(
+            """\
+            </ul>
+            </body>
+            </html>
+            """
+        ))
+
+    # http://repository.example.com/simple/PACKAGE_NAME/
+    for package_name in package_names:
+        _write_package_specific_index(output_dir, package_name, packages[package_name])
 
 
 if __name__ == "__main__":
-    print(main())
+    sys.exit(main(sys.argv[1:]))
diff --git a/requirements.txt b/requirements.txt
@@ -1 +1,2 @@
+packaging==24.2
 pygithub==2.5.0
diff --git a/test_generated_index.sh b/test_generated_index.sh
@@ -1,6 +1,10 @@
 #!/bin/bash
 
 output_dir="$1"
+if [ -z "${output_dir}" ]; then
+  echo "usage: $0 OUTPUT_DIRECTORY" 1>&2
+  exit 1
+fi
 if [ -e "${output_dir}" ]; then
     echo "ERROR: Output directory exists. This script requires that it not exist already." 1>&2
     exit 1
@@ -15,8 +19,9 @@ python3.9 -m venv "${venv_dir}"
 "${venv_dir}/bin/pip" install -r ./requirements.txt
 
 # Generate the Pants PyPi-compatible index.
-mkdir -p "${output_dir}/public/simple"
-"${venv_dir}/bin/python" ./generate_index.py > "${output_dir}/public/simple/index.html"
+"${venv_dir}/bin/python" ./generate_index.py \
+    --url-path-prefix=/simple \
+    "${output_dir}/public/simple"
 
 # Serve a copy of the generated index on port 8080.
 python3.9 -m http.server -d "${output_dir}/public" --bind 127.0.0.1 8080 &
@@ -27,8 +32,8 @@ python3.9 -m http.server -d "${output_dir}/public" --bind 127.0.0.1 8080 &
 pants_venv_dir="${output_dir}/pants-venv"
 python3.9 -m venv "${pants_venv_dir}"
 "${pants_venv_dir}/bin/pip" install -vv \
-  --extra-index-url=http://127.0.0.1:8080/simple/ \
-  pantsbuild.pants==2.18.0a0
+    --extra-index-url=http://127.0.0.1:8080/simple/ \
+    pantsbuild.pants==2.23.0
 
 # Verify that the Pants console script is in the expected location.
 if [ ! -f "${pants_venv_dir}/bin/pants" ]; then