Add branch limit

.github/workflows/cleanup.yaml

@@ -32,5 +32,6 @@ jobs:
         with:
           github_token: ${{ github.token }}
           last_commit_age_days: 100
-          dry_run: no
+          dry_run: yes
           ignore_branches: test_prefix/one,test_prefix/two,test_prefix_ignored/one
+          branch_limit: 1

action.yml

@@ -30,6 +30,10 @@ inputs:
     description: "The API base url to be used in requests to GitHub Enterprise"
     required: false
     default: "https://api.github.com"
+  branch_limit:
+    description: "The max number of branches that can be deleted"
+    required: false
+    default: "100"
 
 outputs:
   deleted_branches: # id of output
@@ -50,5 +54,5 @@ runs:
         python3 main.py --ignore-branches=${{ inputs.ignore_branches }} \
           --last-commit-age-days=${{ inputs.last_commit_age_days }} --allowed-prefixes=${{ inputs.allowed_prefixes }} \
           --dry-run=${{ inputs.dry_run }} --github-token=${{ inputs.github_token }} \
-          --github-base-url=${{ inputs.github_base_url }}
+          --github-base-url=${{ inputs.github_base_url }} --branch-limit=${{ inputs.branch_limit }}
       shell: bash

src/actions.py

@@ -10,7 +10,8 @@ def run_action(options: Options) -> list:
     branches = github.get_deletable_branches(
         last_commit_age_days=options.last_commit_age_days,
         ignore_branches=options.ignore_branches,
-        allowed_prefixes=options.allowed_prefixes
+        allowed_prefixes=options.allowed_prefixes,
+        branch_limit=options.branch_limit
     )
 
     print(f"Branches queued for deletion: {branches}")

src/github.py

@@ -22,8 +22,12 @@ def get_deletable_branches(
             self,
             last_commit_age_days: int,
             ignore_branches: list[str],
-            allowed_prefixes: list[str]
+            allowed_prefixes: list[str],
+            branch_limit: int
     ) -> list[str]:
+        if branch_limit < 1:
+            return []
+
         # Default branch might not be protected
         default_branch = self.get_default_branch()
 
@@ -91,6 +95,10 @@ def get_deletable_branches(
                 print(f'Branch `{branch_name}` meets the criteria for deletion')
                 deletable_branches.append(branch_name)
 
+                # Exit early if we have reached our branch limit
+                if len(deletable_branches) == branch_limit:
+                    return deletable_branches
+
             # Re-request next page
             current_page += 1
 

src/io.py

@@ -13,7 +13,8 @@ def __init__(
             github_token: str,
             github_repo: str,
             dry_run: bool = True,
-            github_base_url: str = DEFAULT_GITHUB_API_URL
+            github_base_url: str = DEFAULT_GITHUB_API_URL,
+            branch_limit: int
     ):
         self.ignore_branches = ignore_branches
         self.last_commit_age_days = last_commit_age_days
@@ -22,6 +23,7 @@ def __init__(
         self.github_repo = github_repo
         self.dry_run = dry_run
         self.github_base_url = github_base_url
+        self.branch_limit = branch_limit
 
 
 class InputParser:
@@ -58,6 +60,13 @@ def get_args() -> argparse.Namespace:
             help="Whether to delete branches at all. Defaults to 'yes'. Possible values: yes, no (case sensitive)"
         )
 
+        parser.add_argument(
+            "--branch-limit",
+            help="The max number of branches that can be deleted",
+            default=100,
+            type=int,
+        )
+
         return parser.parse_args()
 
     def parse_input(self) -> Options:
@@ -83,7 +92,8 @@ def parse_input(self) -> Options:
             dry_run=dry_run,
             github_token=args.github_token,
             github_repo=getenv('GITHUB_REPOSITORY'),
-            github_base_url=args.github_base_url
+            github_base_url=args.github_base_url,
+            branch_limit=args.branch_limit
         )