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,35 +655,94 @@ 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]:
+    return [re.compile(i) for i in exclude]
+
+
+def get_syncer(
+    basedir: Path,
+    local_db: LocalVersionKV,
+    sync_packages: bool,
+    shadowmire_upstream: Optional[str],
+) -> SyncBase:
+    syncer: SyncBase
+    if shadowmire_upstream:
+        syncer = SyncPlainHTTP(
+            upstream=shadowmire_upstream,
             basedir=basedir,
             local_db=local_db,
-                sync_packages=args.sync_packages,
+            sync_packages=sync_packages,
         )
     else:
-            sync = SyncPyPI(
+        syncer = SyncPyPI(
-                basedir=basedir, local_db=local_db, sync_packages=args.sync_packages
+            basedir=basedir, local_db=local_db, sync_packages=sync_packages
         )
+    return syncer
+
+
+@cli.command()
+@click.pass_context
+@sync_shared_args
+def sync(
+    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 = local_db.dump()
-        plan = sync.determine_sync_plan(local, args.excludes)
+    plan = syncer.determine_sync_plan(local, excludes)
     # save plan for debugging
     with overwrite(basedir / "plan.json") as f:
         json.dump(plan, f, default=vars)
-        sync.do_sync_plan(plan)
+    syncer.do_sync_plan(plan)
-        sync.finalize()
+    syncer.finalize()
-    elif args.command == "genlocal":
+
+
+@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
@@ -677,26 +752,27 @@ def main(args: argparse.Namespace) -> None:
     local_db.nuke(commit=False)
     local_db.batch_set(local)
     local_db.dump_json()
-    elif args.command == "verify":
+
-        if args.shadowmire_upstream:
+
-            sync = SyncPlainHTTP(
+@cli.command()
-                upstream=args.shadowmire_upstream,
+@click.pass_context
-                basedir=basedir,
+@sync_shared_args
-                local_db=local_db,
+def verify(
-                sync_packages=args.sync_packages,
+    ctx: click.Context,
-            )
+    sync_packages: bool,
-        else:
+    shadowmire_upstream: Optional[str],
-            sync = SyncPyPI(
+    exclude: tuple[str],
-                basedir=basedir, local_db=local_db, sync_packages=args.sync_packages
+) -> 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(
+    simple_dirs = set([i.name for i in (basedir / "simple").iterdir() if i.is_dir()])
-            [i.name for i in (basedir / "simple").iterdir() if i.is_dir()]
-        )
     for package_name in simple_dirs - local_names:
-            sync.do_remove(package_name)
+        syncer.do_remove(package_name)
-        sync.parallel_update(list(local_names))
+    syncer.parallel_update(list(local_names))
-        sync.finalize()
+    syncer.finalize()
     # clean up unreferenced package files
     ref_set = set()
     for sname in simple_dirs:
@@ -712,47 +788,4 @@ def main(args: argparse.Namespace) -> None:
 
 
 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)