customize args for decorator

Signed-off-by: vsoch <vsoch@users.noreply.github.com>

pull-request.py

@@ -1,9 +1,10 @@
 #!/usr/bin/env python3
 
-import sys
-import os
 import json
+import os
 import requests
+import sys
+import time
 
 ################################################################################
 # Helper Functions
@@ -35,6 +36,28 @@ def check_events_json():
     return events
 
 
+def retry(attempts=5, timeout=30):
+    """
+    A simple retry decorator
+
+    We assume rate limiting is in minutes, so we set timeout to 30.
+    """
+    def decorator(func):
+        def inner(*args, **kwargs):
+            attempt = 0
+            while attempt < attempts:
+                try:
+                    return func(*args, **kwargs)
+                except Exception as e:
+                    sleep = timeout * 2**attempt
+                    print(f"Retrying in {sleep} seconds - error: {e}")
+                    time.sleep(sleep)
+                    attempt += 1
+            return func(*args, **kwargs)
+        return inner
+    return decorator
+
+
 def abort_if_fail(response, reason):
     """If PASS_ON_ERROR, don't exit. Otherwise exit with an error and print
     the reason.
@@ -57,7 +80,8 @@ def abort_if_fail(response, reason):
 
 
 def parse_into_list(values):
-    """A list of reviewers or assignees to parse from a string to a list
+    """
+    A list of reviewers or assignees to parse from a string to a list
 
     Parameters:
     values (str) : a list of space separated, quoted values to parse to a list
@@ -70,7 +94,8 @@ def parse_into_list(values):
 
 
 def set_env_and_output(name, value):
-    """helper function to echo a key/value pair to the environement file
+    """
+    Helper function to echo a key/value pair to the environement file
 
     Parameters:
     name (str)  : the name of the environment variable
@@ -84,8 +109,10 @@ def set_env_and_output(name, value):
             environment_file.write("%s=%s\n" % (name, value))
 
 
+@retry(attempts=5, timeout=60)
 def open_pull_request(title, body, target, source, is_draft=False, can_modify=True):
-    """Open pull request opens a pull request with a given body and content,
+    """
+    Open pull request opens a pull request with a given body and content,
     and sets output variables. An unparsed response is returned.
 
     Parameters:
@@ -116,8 +143,10 @@ def open_pull_request(title, body, target, source, is_draft=False, can_modify=Tr
     return response
 
 
+@retry()
 def update_pull_request(entry, title, body, target, state=None):
-    """Given an existing pull request, update it.
+    """
+    Given an existing pull request, update it.
 
     Parameters:
     entry      (dict) : the pull request metadata
@@ -145,7 +174,8 @@ def update_pull_request(entry, title, body, target, state=None):
 
 
 def set_pull_request_groups(response):
-    """Given a response for an open or updated PR, set metadata
+    """
+    Given a response for an open or updated PR, set metadata
 
     Parameters:
     response (requests.Response) : a requests response, unparsed
@@ -166,6 +196,7 @@ def set_pull_request_groups(response):
     set_env_and_output("PULL_REQUEST_URL", html_url)
 
 
+@retry(attempts=5, timeout=30)
 def list_pull_requests(target, source):
     """Given a target and source, return a list of pull requests that match
     (or simply exit given some kind of error code)
@@ -188,8 +219,10 @@ def list_pull_requests(target, source):
     return response.json()
 
 
+@retry(attempts=5, timeout=30)
 def add_assignees(entry, assignees):
-    """Given a pull request metadata (from create or update) add assignees
+    """
+    Given a pull request metadata (from create or update) add assignees
 
     Parameters:
     entry (dict)    : the pull request metadata
@@ -218,7 +251,8 @@ def add_assignees(entry, assignees):
 
 
 def find_pull_request(listing, source):
-    """Given a listing and a source, find a pull request based on the source
+    """
+    Given a listing and a source, find a pull request based on the source
     (the branch name).
 
     Parameters:
@@ -232,8 +266,11 @@ def find_pull_request(listing, source):
                 return entry
 
 
+@retry(attempts=5, timeout=30)
 def find_default_branch():
-    """Find default branch for a repo (only called if branch not provided)"""
+    """
+    Find default branch for a repo (only called if branch not provided)
+    """
     response = requests.get(REPO_URL)
 
     # Case 1: 401, 404 might need a token
@@ -247,8 +284,10 @@ def find_default_branch():
     return default_branch
 
 
+@retry(attempts=5, timeout=30)
 def add_reviewers(entry, reviewers, team_reviewers):
-    """Given regular or team reviewers, add them to a PR.
+    """
+    Given regular or team reviewers, add them to a PR.
 
     Parameters:
     entry (dict) : the pull request metadata
@@ -357,7 +396,10 @@ def main():
         print("No branch prefix is set, all branches will be used.")
 
     # Default to project default branch if none provided
-    pull_request_branch = os.environ.get("PULL_REQUEST_BRANCH", find_default_branch())
+    pull_request_branch = os.environ.get("PULL_REQUEST_BRANCH")
+    if not pull_request_branch:
+        pull_request_branch = find_default_branch()
+
     print("Pull requests will go to %s" % pull_request_branch)
 
     # Pull request draft
@@ -431,7 +473,6 @@ def main():
 
     # If the prefix for the branch matches
     if not branch_prefix or from_branch.startswith(branch_prefix):
-
         # Pull request body (optional)
         pull_request_body = os.environ.get(
             "PULL_REQUEST_BODY",