From 046775c0f7eb1891f850f6f9b970974df74a4ac3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torjus=20H=C3=A5kestad?= <torjus@usit.uio.no>
Date: Tue, 1 Oct 2024 01:56:13 +0200
Subject: [PATCH] Add option for checking specific branch

---
 nixprstatus/__main__.py | 53 +++++++++++++++++++++++++++++++----------
 1 file changed, 40 insertions(+), 13 deletions(-)

diff --git a/nixprstatus/__main__.py b/nixprstatus/__main__.py
index 010e0d2..e336ec7 100644
--- a/nixprstatus/__main__.py
+++ b/nixprstatus/__main__.py
@@ -1,5 +1,6 @@
 import typer
 import requests
+from typing import Annotated
 from rich.console import Console
 
 app = typer.Typer()
@@ -11,10 +12,25 @@ DEFAULT_HEADERS = {
 BRANCHES = ["nixos-unstable-small", "nixos-unstable", "nixos-24.05"]
 
 
+def commit_in_branch(commit_sha: str, branch: str) -> bool:
+    url = f"https://api.github.com/repos/NixOS/nixpkgs/compare/{branch}...{commit_sha}"
+    commit_response = requests.get(url, headers=DEFAULT_HEADERS)
+    commit_response.raise_for_status()
+    status = commit_response.json().get("status")
+
+    if status in ["identical", "behind"]:
+        return True
+    return False
+
+
 @app.command()
-def pr(pr: str):
+def pr(
+    pr: str,
+    branch: Annotated[str | None, typer.Option(help="Check specific branch")] = None,
+):
     """Get status of pull request"""
     url = f"https://api.github.com/repos/NixOS/nixpkgs/pulls/{pr}"
+    typer.echo(f"branch {branch}")
 
     pr_response = requests.get(url, headers=DEFAULT_HEADERS)
     pr_response.raise_for_status()
@@ -23,7 +39,24 @@ def pr(pr: str):
 
     console = Console()
 
-    if pr_data["merged"] is False:
+    merged = pr_data["merged"]
+
+    if branch and not merged:
+        console.print(f":x: {branch}")
+        return
+
+    commit_sha = pr_data.get("merge_commit_sha")
+
+    # Check only specified branch
+    if branch:
+        in_branch = commit_in_branch(commit_sha, branch)
+        if in_branch:
+            console.print(f":white_check_mark: {branch}", highlight=False)
+        else:
+            console.print(f":x: {branch}", highlight=False)
+        return
+
+    if not merged:
         console.print(":x: master", highlight=False)
         for branch in BRANCHES:
             console.print(f":x: {branch}", highlight=False)
@@ -31,18 +64,12 @@ def pr(pr: str):
 
     console.print(":white_check_mark: master", highlight=False)
 
-    commit_sha = pr_data["merge_commit_sha"]
-
-    for branch in BRANCHES:
-        url = f"https://api.github.com/repos/NixOS/nixpkgs/compare/{branch}...{commit_sha}"
-        commit_response = requests.get(url, headers=DEFAULT_HEADERS)
-        commit_response.raise_for_status()
-        status = commit_response.json().get("status")
-
-        if status in ["identical", "behind"]:
-            console.print(f":white_check_mark: {branch}", highlight=False)
+    for b in BRANCHES:
+        in_branch = commit_in_branch(commit_sha, b)
+        if in_branch:
+            console.print(f":white_check_mark: {b}", highlight=False)
         else:
-            console.print(f":x: {branch}", highlight=False)
+            console.print(f":x: {b}", highlight=False)
 
 
 def main():