Use click to replace argparse

Make it simpler and easier to maintain command line arguments.

requirements.txt

@@ -1,2 +1,3 @@
 requests==2.32.3
 tqdm==4.66.4
+click==8.1.7

shadowmire.py

@@ -10,18 +10,19 @@ from pathlib import Path
 from html.parser import HTMLParser
 import logging
 import html
-import argparse
 import os
 from contextlib import contextmanager
 import sqlite3
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import signal
 import requests
+import click
 from tqdm import tqdm
 from requests.adapters import HTTPAdapter, Retry
 
 logger = logging.getLogger(__name__)
 
+
 USER_AGENT = "Shadowmire (https://github.com/taoky/shadowmire)"
 
 # Note that it's suggested to use only 3 workers for PyPI.
@@ -32,6 +33,21 @@ if WORKERS > 10:
     )
     logger.warning("Don't blame me if you were banned!")
 
+# https://github.com/pypa/bandersnatch/blob/a05af547f8d1958217ef0dc0028890b1839e6116/src/bandersnatch_filter_plugins/prerelease_name.py#L18C1-L23C6
+PRERELEASE_PATTERNS = (
+    re.compile(r".+rc\d+$"),
+    re.compile(r".+a(lpha)?\d+$"),
+    re.compile(r".+b(eta)?\d+$"),
+    re.compile(r".+dev\d+$"),
+)
+
+
+def is_version_prerelease(version: str) -> bool:
+    for p in PRERELEASE_PATTERNS:
+        if p.match(version):
+            return True
+    return False
+
 
 class PackageNotFoundError(Exception):
     pass
@@ -195,8 +211,8 @@ class CustomXMLRPCTransport(xmlrpc.client.Transport):
 def create_requests_session() -> requests.Session:
     s = requests.Session()
     retries = Retry(total=3, backoff_factor=0.1)
-    s.mount('http://', HTTPAdapter(max_retries=retries))
+    s.mount("http://", HTTPAdapter(max_retries=retries))
-    s.mount('https://', HTTPAdapter(max_retries=retries))
+    s.mount("https://", HTTPAdapter(max_retries=retries))
     s.headers.update({"User-Agent": USER_AGENT})
     return s
 
@@ -639,120 +655,137 @@ def get_local_serial(package_simple_path: Path) -> Optional[int]:
         return None
 
 
-def main(args: argparse.Namespace) -> None:
+def sync_shared_args(func):
+    shared_options = [
+        click.option(
+            "--sync-packages",
+            is_flag=True,
+            help="Sync packages instead of just indexes",
+        ),
+        click.option(
+            "--shadowmire-upstream",
+            required=False,
+            type=str,
+            help="Use another upstream using shadowmire instead of PyPI",
+        ),
+        click.option(
+            "--exclude", multiple=True, help="Remote package names to exclude. Regex."
+        ),
+    ]
+    for option in shared_options[::-1]:
+        func = option(func)
+    return func
+
+
+@click.group()
+@click.pass_context
+def cli(ctx: click.Context) -> None:
     log_level = logging.DEBUG if os.environ.get("DEBUG") else logging.INFO
     logging.basicConfig(level=log_level)
-    logger.debug(args)
+    ctx.ensure_object(dict)
 
     basedir = Path(os.environ.get("REPO", "."))
     local_db = LocalVersionKV(basedir / "local.db", basedir / "local.json")
 
-    sync: SyncBase
+    ctx.obj["basedir"] = basedir
-    if args.command == "sync":
+    ctx.obj["local_db"] = local_db
-        if args.shadowmire_upstream:
+
-            sync = SyncPlainHTTP(
+
-                upstream=args.shadowmire_upstream,
+def exclude_to_excludes(exclude: tuple[str]) -> list[re.Pattern]:
-                basedir=basedir,
+    return [re.compile(i) for i in exclude]
-                local_db=local_db,
+
-                sync_packages=args.sync_packages,
+
-            )
+def get_syncer(
-        else:
+    basedir: Path,
-            sync = SyncPyPI(
+    local_db: LocalVersionKV,
-                basedir=basedir, local_db=local_db, sync_packages=args.sync_packages
+    sync_packages: bool,
-            )
+    shadowmire_upstream: Optional[str],
-        local = local_db.dump()
+) -> SyncBase:
-        plan = sync.determine_sync_plan(local, args.excludes)
+    syncer: SyncBase
-        # save plan for debugging
+    if shadowmire_upstream:
-        with overwrite(basedir / "plan.json") as f:
+        syncer = SyncPlainHTTP(
-            json.dump(plan, f, default=vars)
+            upstream=shadowmire_upstream,
-        sync.do_sync_plan(plan)
+            basedir=basedir,
-        sync.finalize()
+            local_db=local_db,
-    elif args.command == "genlocal":
+            sync_packages=sync_packages,
-        local = {}
-        for package_path in (basedir / "simple").iterdir():
-            package_name = package_path.name
-            serial = get_local_serial(package_path)
-            if serial:
-                local[package_name] = serial
-        local_db.nuke(commit=False)
-        local_db.batch_set(local)
-        local_db.dump_json()
-    elif args.command == "verify":
-        if args.shadowmire_upstream:
-            sync = SyncPlainHTTP(
-                upstream=args.shadowmire_upstream,
-                basedir=basedir,
-                local_db=local_db,
-                sync_packages=args.sync_packages,
-            )
-        else:
-            sync = SyncPyPI(
-                basedir=basedir, local_db=local_db, sync_packages=args.sync_packages
-            )
-        local_names = set(local_db.keys())
-        simple_dirs = set(
-            [i.name for i in (basedir / "simple").iterdir() if i.is_dir()]
         )
-        for package_name in simple_dirs - local_names:
+    else:
-            sync.do_remove(package_name)
+        syncer = SyncPyPI(
-        sync.parallel_update(list(local_names))
+            basedir=basedir, local_db=local_db, sync_packages=sync_packages
-        sync.finalize()
+        )
-        # clean up unreferenced package files
+    return syncer
-        ref_set = set()
+
-        for sname in simple_dirs:
+
-            sd = basedir / "simple" / sname
+@cli.command()
-            hrefs = get_existing_hrefs(sd)
+@click.pass_context
-            for i in hrefs:
+@sync_shared_args
-                ref_set.add(str((sd / i).resolve()))
+def sync(
-        for file in (basedir / "packages").glob("*/*/*/*"):
+    ctx: click.Context,
-            file = file.resolve()
+    sync_packages: bool,
-            if str(file) not in ref_set:
+    shadowmire_upstream: Optional[str],
-                logger.info("removing unreferenced %s", file)
+    exclude: tuple[str],
-                file.unlink()
+) -> None:
+    basedir = ctx.obj["basedir"]
+    local_db = ctx.obj["local_db"]
+    excludes = exclude_to_excludes(exclude)
+    syncer = get_syncer(basedir, local_db, sync_packages, shadowmire_upstream)
+    local = local_db.dump()
+    plan = syncer.determine_sync_plan(local, excludes)
+    # save plan for debugging
+    with overwrite(basedir / "plan.json") as f:
+        json.dump(plan, f, default=vars)
+    syncer.do_sync_plan(plan)
+    syncer.finalize()
+
+
+@cli.command()
+@click.pass_context
+def genlocal(ctx: click.Context) -> None:
+    basedir = ctx.obj["basedir"]
+    local_db = ctx.obj["local_db"]
+    local = {}
+    for package_path in (basedir / "simple").iterdir():
+        package_name = package_path.name
+        serial = get_local_serial(package_path)
+        if serial:
+            local[package_name] = serial
+    local_db.nuke(commit=False)
+    local_db.batch_set(local)
+    local_db.dump_json()
+
+
+@cli.command()
+@click.pass_context
+@sync_shared_args
+def verify(
+    ctx: click.Context,
+    sync_packages: bool,
+    shadowmire_upstream: Optional[str],
+    exclude: tuple[str],
+) -> None:
+    basedir = ctx.obj["basedir"]
+    local_db = ctx.obj["local_db"]
+    excludes = exclude_to_excludes(exclude)
+    syncer = get_syncer(basedir, local_db, sync_packages, shadowmire_upstream)
+    local_names = set(local_db.keys())
+    simple_dirs = set([i.name for i in (basedir / "simple").iterdir() if i.is_dir()])
+    for package_name in simple_dirs - local_names:
+        syncer.do_remove(package_name)
+    syncer.parallel_update(list(local_names))
+    syncer.finalize()
+    # clean up unreferenced package files
+    ref_set = set()
+    for sname in simple_dirs:
+        sd = basedir / "simple" / sname
+        hrefs = get_existing_hrefs(sd)
+        for i in hrefs:
+            ref_set.add(str((sd / i).resolve()))
+    for file in (basedir / "packages").glob("*/*/*/*"):
+        file = file.resolve()
+        if str(file) not in ref_set:
+            logger.info("removing unreferenced %s", file)
+            file.unlink()
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser("shadowmire: lightweight PyPI syncing tool")
+    cli(obj={})
-    subparsers = parser.add_subparsers(dest="command")
-
-    parser_sync = subparsers.add_parser("sync", help="Sync from upstream")
-    parser_sync.add_argument(
-        "--sync-packages",
-        help="Sync packages instead of just indexes",
-        action="store_true",
-    )
-    parser_sync.add_argument(
-        "--exclude", help="Remote package names to exclude. Regex.", nargs="*"
-    )
-    parser_sync.add_argument(
-        "--shadowmire-upstream",
-        help="Use another upstream using shadowmire instead of PyPI",
-        type=str,
-    )
-    parser_genlocal = subparsers.add_parser(
-        "genlocal", help="(Re)generate local db and json from simple/"
-    )
-    parser_verify = subparsers.add_parser(
-        "verify",
-        help="Verify existing sync from local db, download missing things, remove unreferenced packages",
-    )
-    parser_verify.add_argument(
-        "--sync-packages",
-        help="Sync packages instead of just indexes",
-        action="store_true",
-    )
-    parser_verify.add_argument(
-        "--shadowmire-upstream",
-        help="Use another upstream using shadowmire instead of PyPI",
-        type=str,
-    )
-
-    args = parser.parse_args()
-    if args.command is None:
-        parser.print_help()
-        sys.exit(1)
-    if args.command == "sync" and args.exclude:
-        args.excludes = [re.compile(i) for i in args.exclude]
-    else:
-        args.excludes = []
-    main(args)