customize args for decorator

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

CHANGELOG.md

@@ -14,6 +14,7 @@ represented by the pull requests that fixed them. Critical items to know are:
 Versions correspond with GitHub releases that can be referenced with @ using actions.
 
 ## [master](https://github.com/vsoch/pull-request-action/tree/master) (master)
+  - bugfix of missing output values (1.0.23)
   - bugfix of token handling if 401 error received (missing 401 case) (1.0.21)
   - bugfix of writing to environment file (missing newline) (1.0.19)
   - bugfix of missing from branch with scheduled run (1.0.16)

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
@@ -69,21 +93,26 @@ def parse_into_list(values):
     return [x.strip() for x in values.split(" ")]
 
 
-def set_env(name, value):
-    """helper function to echo a key/value pair to the environement file
+def set_env_and_output(name, value):
+    """
+    Helper function to echo a key/value pair to the environement file
 
     Parameters:
     name (str)  : the name of the environment variable
     value (str) : the value to write to file
     """
-    environment_file_path = os.environ.get("GITHUB_ENV")
+    for env_var in ("GITHUB_ENV", "GITHUB_OUTPUT"):
+        environment_file_path = os.environ.get(env_var)
+        print("Writing %s=%s to %s" % (name, value, env_var))
 
-    with open(environment_file_path, "a") as environment_file:
-        environment_file.write("%s=%s\n" % (name, value))
+        with open(environment_file_path, "a") as environment_file:
+            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:
@@ -114,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
@@ -143,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
@@ -159,14 +191,12 @@ def set_pull_request_groups(response):
     number = response.get("number")
     html_url = response.get("html_url")
     print("Number opened for PR is %s" % number)
-    set_env("PULL_REQUEST_NUMBER", number)
-    print("pull_request_number=%s >> $GITHUB_OUTPUT" % number)
-    set_env("PULL_REQUEST_RETURN_CODE", pull_request_return_code)
-    print("pull_request_return_code=%s >> $GITHUB_OUTPUT" % pull_request_return_code)
-    set_env("PULL_REQUEST_URL", html_url)
-    print("pull_request_url=%s >> $GITHUB_OUTPUT" % html_url)
+    set_env_and_output("PULL_REQUEST_NUMBER", number)
+    set_env_and_output("PULL_REQUEST_RETURN_CODE", pull_request_return_code)
+    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)
@@ -189,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
@@ -215,12 +247,12 @@ def add_assignees(entry, assignees):
     print("::group::github assignees response")
     print(response.json())
     print("::endgroup::github assignees response")
-    set_env("ASSIGNEES_RETURN_CODE", assignees_return_code)
-    print("assignees_return_code=%s >> $GITHUB_OUTPUT" % assignees_return_code)
+    set_env_and_output("ASSIGNEES_RETURN_CODE", assignees_return_code)
 
 
 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:
@@ -234,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
@@ -249,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
@@ -272,9 +309,7 @@ def add_reviewers(entry, reviewers, team_reviewers):
     print("::group::github reviewers response")
     print(response.json())
     print("::endgroup::github reviewers response")
-    set_env("REVIEWERS_RETURN_CODE", reviewers_return_code)
-    print("reviewers_return_code=%s >> $GITHUB_OUTPUT" % reviewers_return_code)
-    print("Add reviewers return code: %s" % reviewers_return_code)
+    set_env_and_output("REVIEWERS_RETURN_CODE", reviewers_return_code)
 
 
 ################################################################################
@@ -361,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
@@ -435,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",