Only support pypi in cleanup script

Author: evertlammerts

.github/workflows/cleanup_pypi.yml

@@ -34,11 +34,6 @@ jobs:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-
-      - if: ${{ vars.PYPI_HOST == '' }}
-        run: |
-          echo "Error: PYPI_HOST is not set in CI environment '${{ inputs.environment }}'"
-          exit 1
       - if: ${{ vars.PYPI_CLEANUP_USERNAME == '' }}
         run: |
           echo "Error: PYPI_CLEANUP_USERNAME is not set in CI environment '${{ inputs.environment }}'"
@@ -56,9 +51,10 @@ jobs:
       - name: Run Cleanup
         run: |
           set -x
+          pypi_index_flag=${{ inputs.environment == 'production.pypi' && '--prod' || '--test' }}
           uv sync --only-group pypi --no-install-project
           uv run --no-sync -s scripts/pypi_cleanup.py ${{ inputs.dry-run && '--dry' || '' }} \
-            --index-hostname "${{ vars.PYPI_HOST }}" \
+            ${pypi_index_flag} \
             --username "${{ vars.PYPI_CLEANUP_USERNAME }}" \
             --max-nightlies ${{ vars.PYPI_MAX_NIGHTLIES }} > cleanup_output 2>&1
 

scripts/pypi_cleanup.py

@@ -14,10 +14,6 @@
 import requests
 from requests.exceptions import RequestException
 
-import argparse
-import re
-import os
-
 def valid_hostname(hostname):
     """Validate hostname format"""
     if len(hostname) > 253:
@@ -41,7 +37,9 @@ def non_empty_string(value):
     epilog="Environment variables required (unless --dry): PYPI_CLEANUP_PASSWORD, PYPI_CLEANUP_OTP"
 )
 parser.add_argument("--dry", action="store_true", help="Show what would be deleted but don't actually do it")
-parser.add_argument("-i", "--index-hostname", type=valid_hostname, required=True, help="Index hostname (required)")
+host_group = parser.add_mutually_exclusive_group(required=True)
+host_group.add_argument("--prod", action="store_true", help="Use production PyPI (pypi.org)")
+host_group.add_argument("--test", action="store_true", help="Use test PyPI (test.pypi.org)")
 parser.add_argument("-m", "--max-nightlies", type=int, default=2, help="Max number of nightlies of unreleased versions (default=2)")
 parser.add_argument("-u", "--username", type=non_empty_string, help="Username (required unless --dry)")
 args = parser.parse_args()
@@ -62,21 +60,25 @@ def non_empty_string(value):
     if not otp:
         parser.error("PYPI_CLEANUP_OTP environment variable is required when not in dry-run mode")
 
-print(f"Dry run: {args.dry}")
-print(f"Max nightlies: {args.max_nightlies}")
+dry_run = args.dry
+pypi_username = args.username
+max_dev_releases = args.max_nightlies
+host = None
+if args.prod:
+    host = 'pypi.org'
+elif args.test:
+    host = 'test.pypi.org'
+pypi_url = 'https://{}/'.format(host)
+pypi_password = password
+pypi_otp = otp
+
+print(f"Dry run: {dry_run}")
+print(f"Max nightlies: {max_dev_releases}")
 if not args.dry:
-    print(f"Hostname: {args.index_hostname}")
-    print(f"Username: {args.username}")
+    print(f"URL: {pypi_url}")
+    print(f"Username: {pypi_username}")
     print("Password and OTP loaded from environment variables")
 
-# deletes old dev wheels from pypi. evil hack.
-actually_delete = not args.dry
-pypi_username = args.username or "user"
-max_dev_releases = args.max_nightlies
-host = 'https://{}/'.format(args.index_hostname)
-pypi_password = password or "password"
-pypi_otp = otp or "otp"
-
 patterns = [re.compile(r".*\.dev\d+$")]
 ###### NOTE: This code is taken from the pypi-cleanup package (https://github.com/arcivanov/pypi-cleanup/tree/master)
 class CsfrParser(HTMLParser):
@@ -300,4 +302,4 @@ def run(self):
                     logging.info(f"Would be deleting {self.package} version {pkg_ver}, but not doing it!")
 
 
-PypiCleanup(host, pypi_username, 'duckdb', pypi_password, pypi_otp, patterns, actually_delete, max_dev_releases).run()
+PypiCleanup(pypi_url, pypi_username, 'duckdb', pypi_password, pypi_otp, patterns, not dry_run, max_dev_releases).run()