Add retry on rate limits to publish (#21)

Adds retry when a 429 is encountered and better error display in the
`publish` CLI.

Unfortunately the rate limit for new projects looks relatively severe at
20 packages / hr

https://github.com/pypi/warehouse/blob/d858a996fe543131931955d8e8cc96b8aa7261a1/warehouse/config.py#L352-L369

Also adds some mocked twine test coverage to `publish`

src/packse/cli.py

@@ -4,7 +4,12 @@
 from pathlib import Path
 
 from packse.build import build
-from packse.error import BuildError, DestinationAlreadyExists, UserError
+from packse.error import (
+    BuildError,
+    DestinationAlreadyExists,
+    PublishError,
+    UserError,
+)
 from packse.publish import publish
 from packse.view import view
 
@@ -37,6 +42,9 @@ def entrypoint():
     except BuildError as exc:
         print(f"{exc}", file=sys.stderr)
         exit(1)
+    except PublishError as exc:
+        print(f"{exc}", file=sys.stderr)
+        exit(1)
 
 
 def _call_build(args):
@@ -48,7 +56,12 @@ def _call_view(args):
 
 
 def _call_publish(args):
-    publish(args.targets, dry_run=args.dry_run, skip_existing=args.skip_existing)
+    publish(
+        args.targets,
+        dry_run=args.dry_run,
+        skip_existing=args.skip_existing,
+        retry_on_rate_limit=args.retry,
+    )
 
 
 def _root_parser():
@@ -95,6 +108,11 @@ def _add_publish_parser(subparsers):
         action="store_true",
         help="Skip existing distributions instead of failing.",
     )
+    parser.add_argument(
+        "--retry",
+        action="store_true",
+        help="Retry when rate limits are encountered instead of failing. Note retries must be very long to succeed.",
+    )
     _add_shared_arguments(parser)
 
 

src/packse/error.py

@@ -35,14 +35,24 @@ def __init__(self, message, output) -> None:
 
 
 class PublishError(PackseError):
+    pass
+
+
+class PublishToolError(PublishError):
     def __init__(self, message, output) -> None:
         message = message + ":\n" + textwrap.indent(output, " " * 4)
         super().__init__(message)
 
 
-class PublishAlreadyExists(PackseError):
+class PublishAlreadyExists(PublishError):
+    def __init__(self, package: str) -> None:
+        message = f"Publish for '{package}' already exists"
+        super().__init__(message)
+
+
+class PublishRateLimit(PublishError):
     def __init__(self, package: str) -> None:
-        message = f"Publish for '{package}' already exists."
+        message = f"Publish of '{package}' failed due to rate limits"
         super().__init__(message)
 
 

src/packse/publish.py

@@ -10,30 +10,70 @@
 from packse.error import (
     InvalidPublishTarget,
     PublishAlreadyExists,
-    PublishError,
+    PublishRateLimit,
+    PublishToolError,
 )
 
 logger = logging.getLogger(__name__)
 
 
-def publish(targets: list[Path], skip_existing: bool, dry_run: bool):
+# Rate limits are defined at https://github.com/pypi/warehouse/blob/d858a996fe543131931955d8e8cc96b8aa7261a1/warehouse/config.py#L352-L369
+# At the time of writing, rate limits enforce a maximum of 20 new projects / hour.
+# Generally, retries aren't going to help.
+MAX_ATTEMPTS = 3
+RETRY_TIME = 60 * 2  # Start with two minutes
+RETRY_BACKOFF_FACTOR = 1.5
+
+
+def publish(
+    targets: list[Path], skip_existing: bool, dry_run: bool, retry_on_rate_limit: bool
+):
     for target in targets:
         if not target.is_dir():
             raise InvalidPublishTarget(target, reason="Not a directory.")
 
     s = "" if len(targets) == 1 else "s"
     logger.info("Publishing %s target%s...", len(targets), s)
-    for target in targets:
+    for target in sorted(targets):
         logger.info("Publishing '%s'...", target.name)
         for distfile in sorted(target.iterdir()):
-            try:
-                publish_package_distribution(distfile, dry_run)
-            except PublishAlreadyExists:
-                if not skip_existing:
-                    raise
-                logger.info("Skipped '%s': already published.", distfile.name)
-            else:
-                logger.info("Published '%s'", distfile.name)
+            publish_package_distribution_with_retries(
+                distfile,
+                skip_existing=skip_existing,
+                dry_run=dry_run,
+                max_attempts=MAX_ATTEMPTS if retry_on_rate_limit else 1,
+            )
+
+
+def publish_package_distribution_with_retries(
+    target: Path,
+    skip_existing: bool,
+    dry_run: bool,
+    max_attempts: int,
+):
+    attempts = 0
+    retry_time = RETRY_TIME
+    while attempts < max_attempts:
+        retry_time = retry_time * RETRY_BACKOFF_FACTOR
+        attempts += 1
+        try:
+            publish_package_distribution(target, dry_run)
+        except PublishAlreadyExists:
+            if not skip_existing:
+                raise
+            logger.info("Skipped '%s': already published.", target.name)
+        except PublishRateLimit:
+            if attempts >= max_attempts:
+                raise
+            logger.warning(
+                "Encountered rate limit publishing '%s', retrying in ~%d minutes",
+                target.name,
+                retry_time // 60,
+            )
+            time.sleep(retry_time)
+        else:
+            logger.info("Published '%s'", target.name)
+            break
 
 
 def publish_package_distribution(target: Path, dry_run: bool) -> None:
@@ -47,13 +87,19 @@ def publish_package_distribution(target: Path, dry_run: bool) -> None:
 
     start_time = time.time()
     try:
-        output = subprocess.check_output(command, stderr=subprocess.STDOUT)
+        import os
+
+        output = subprocess.check_output(
+            command, stderr=subprocess.STDOUT, env=os.environ
+        )
     except subprocess.CalledProcessError as exc:
         output = exc.output.decode()
         if "File already exists" in output:
             raise PublishAlreadyExists(target.name)
-        raise PublishError(
-            f"Publishing {target} with twine failed",
+        if "HTTPError: 429 Too Many Requests" in output:
+            raise PublishRateLimit(target.name)
+        raise PublishToolError(
+            f"Publishing {target.name} with twine failed",
             output,
         )
     else:

tests/__snapshots__/test_publish.ambr

@@ -1,5 +1,5 @@
 # serializer version: 1
-# name: test_publish_example
+# name: test_publish_example_dry_run
   dict({
     'exit_code': 0,
     'stderr': '''
@@ -13,12 +13,86 @@
 
     ''',
     'stdout': '''
-      Would execute: twine upload -r testpypi [PWD]/dist/example-0611cb74/example_0611cb74-0.0.0.tar.gz
-      Would execute: twine upload -r testpypi [PWD]/dist/example-0611cb74/example_0611cb74_a-1.0.0-py3-none-any.whl
-      Would execute: twine upload -r testpypi [PWD]/dist/example-0611cb74/example_0611cb74_a-1.0.0.tar.gz
-      Would execute: twine upload -r testpypi [PWD]/dist/example-0611cb74/example_0611cb74_b-1.0.0-py3-none-any.whl
-      Would execute: twine upload -r testpypi [PWD]/dist/example-0611cb74/example_0611cb74_b-1.0.0.tar.gz
+      Would execute: twine upload -r testpypi [DISTDIR]/example_0611cb74-0.0.0.tar.gz
+      Would execute: twine upload -r testpypi [DISTDIR]/example_0611cb74_a-1.0.0-py3-none-any.whl
+      Would execute: twine upload -r testpypi [DISTDIR]/example_0611cb74_a-1.0.0.tar.gz
+      Would execute: twine upload -r testpypi [DISTDIR]/example_0611cb74_b-1.0.0-py3-none-any.whl
+      Would execute: twine upload -r testpypi [DISTDIR]/example_0611cb74_b-1.0.0.tar.gz
 
     ''',
   })
 # ---
+# name: test_publish_example_twine_fails_with_already_exists
+  dict({
+    'exit_code': 1,
+    'stderr': '''
+      INFO:packse.publish:Publishing 1 target...
+      INFO:packse.publish:Publishing 'example-0611cb74'...
+      Publish for 'example_0611cb74-0.0.0.tar.gz' already exists
+
+    ''',
+    'stdout': '',
+  })
+# ---
+# name: test_publish_example_twine_fails_with_rate_limit
+  dict({
+    'exit_code': 1,
+    'stderr': '''
+      INFO:packse.publish:Publishing 1 target...
+      INFO:packse.publish:Publishing 'example-0611cb74'...
+      Publish of 'example_0611cb74-0.0.0.tar.gz' failed due to rate limits
+
+    ''',
+    'stdout': '',
+  })
+# ---
+# name: test_publish_example_twine_fails_with_unknown_error
+  dict({
+    'exit_code': 1,
+    'stderr': '''
+      INFO:packse.publish:Publishing 1 target...
+      INFO:packse.publish:Publishing 'example-0611cb74'...
+      Publishing example_0611cb74-0.0.0.tar.gz with twine failed:
+          <twine error message>
+
+
+    ''',
+    'stdout': '',
+  })
+# ---
+# name: test_publish_example_twine_succeeds
+  dict({
+    'exit_code': 0,
+    'stderr': '''
+      INFO:packse.publish:Publishing 1 target...
+      INFO:packse.publish:Publishing 'example-0611cb74'...
+      DEBUG:packse.publish:Published example_0611cb74-0.0.0.tar.gz in [TIME]:
+
+          <twine happy message>
+
+      INFO:packse.publish:Published 'example_0611cb74-0.0.0.tar.gz'
+      DEBUG:packse.publish:Published example_0611cb74_a-1.0.0-py3-none-any.whl in [TIME]:
+
+          <twine happy message>
+
+      INFO:packse.publish:Published 'example_0611cb74_a-1.0.0-py3-none-any.whl'
+      DEBUG:packse.publish:Published example_0611cb74_a-1.0.0.tar.gz in [TIME]:
+
+          <twine happy message>
+
+      INFO:packse.publish:Published 'example_0611cb74_a-1.0.0.tar.gz'
+      DEBUG:packse.publish:Published example_0611cb74_b-1.0.0-py3-none-any.whl in [TIME]:
+
+          <twine happy message>
+
+      INFO:packse.publish:Published 'example_0611cb74_b-1.0.0-py3-none-any.whl'
+      DEBUG:packse.publish:Published example_0611cb74_b-1.0.0.tar.gz in [TIME]:
+
+          <twine happy message>
+
+      INFO:packse.publish:Published 'example_0611cb74_b-1.0.0.tar.gz'
+
+    ''',
+    'stdout': '',
+  })
+# ---

tests/common.py

@@ -24,6 +24,7 @@ def snapshot_command(
     snapshot_filesystem: bool = False,
     stderr: bool = True,
     stdout: bool = True,
+    extra_filters: list[tuple[str, str]] | None = None,
 ) -> dict:
     # By default, filter out absolute references to the working directory
     filters = [
@@ -34,12 +35,15 @@ def snapshot_command(
             "[TIME]",
         ),
     ]
+    if extra_filters:
+        filters += extra_filters
 
     process = subprocess.run(
         ["packse"] + command,
         cwd=working_directory,
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
+        env=os.environ,
     )
     result = {
         "exit_code": process.returncode,