Simplify printing

Reviewed-on: #15

README.md

@@ -5,15 +5,27 @@ Nixpkgs PR status checker.
 ## Example
 
 ```console
-$ nixprstatus 345501                                              
+$ nixprstatus pr 345501         
+ktailctl: 0.18.0 -> 0.18.1
 
+✅ merged
 ✅ master
-❌ nixos-unstable-small
+✅ nixos-unstable-small
-❌ nixos-unstable
+✅ nixos-unstable
 ❌ nixos-24.05
+
+$ nixprstatus --help
+Usage: python -m nixprstatus [OPTIONS] COMMAND [ARGS]...
+
+Options:
+  --install-completion  Install completion for the current shell.
+  --show-completion     Show completion for the current shell, to copy it or
+                        customize the installation.
+  --help                Show this message and exit.
+
+Commands:
+  pr         Get merge status of pull request.
+  since      Return the count of commits that has happened between the...
+  watchlist  Manage watchlist.
 ```
 
-## TODO
-
-* Support backported commits
-* JSON output

flake.nix

@@ -32,7 +32,12 @@
       {
         packages = {
           nixprstatus = mkPoetryApplication {
-            projectDir = ./.;
+            projectDir = pkgs.lib.sourceFilesBySuffices ./. [
+              "pyproject.toml"
+              "poetry.lock"
+              "README.md"
+              ".py"
+            ];
             python = pkgs.python312;
             nativeBuildInputs = [ pkgs.installShellFiles ];
             postInstall = ''

nixprstatus/__main__.py

@@ -1,11 +1,14 @@
 import typer
 import json
 from typing import Annotated
-from rich.console import Console
 from nixprstatus.pr import pr_merge_status
 from nixprstatus.pr import commits_since
+from nixprstatus.watchlist import Watchlist
+from nixprstatus.watchlist import OutputFormat
 
-app = typer.Typer()
+app = typer.Typer(rich_markup_mode=None)
+watchlist_app = typer.Typer()
+app.add_typer(watchlist_app, name="watchlist", help="Manage watchlist.")
 
 DEFAULT_HEADERS = {
     "Accept": "application/vnd.github.text+json",
@@ -20,24 +23,17 @@ def pr(
     branches: Annotated[
         list[str] | None, typer.Option(help="Check specific branch")
     ] = None,
+    format: Annotated[
+        OutputFormat, typer.Option(help="Output format")
+    ] = OutputFormat.CONSOLE,
 ):
-    """Get status of pull request"""
+    """Get merge status of pull request."""
-    console = Console()
-
     if branches:
         status = pr_merge_status(pr, branches)
     else:
         status = pr_merge_status(pr)
-    merged = ":white_check_mark: merged" if status.merged else ":x: merged"
-    console.print(merged, highlight=False)
 
-    for branch in status.branches:
+    status.print(format=format)
-        output = (
-            f":white_check_mark: {branch}"
-            if status.branches[branch]
-            else f":x: {branch}"
-        )
-        console.print(output, highlight=False)
 
 
 @app.command()
@@ -60,6 +56,34 @@ def since(
     typer.echo(count)
 
 
+@watchlist_app.command()
+def list(watchlist: str | None = None, format: OutputFormat = OutputFormat.CONSOLE):
+    """List PRs in watchlist."""
+    wl = Watchlist.from_file(path=watchlist)
+    wl.print(format=format)
+
+
+@watchlist_app.command()
+def add(pr: int, watchlist: str | None = None):
+    """Add PR to watchlist."""
+    wl = Watchlist.from_file(path=watchlist)
+    info = wl.add_pr(pr)
+    wl.to_file(path=watchlist)
+    print(f"Added #{info.pr}: {info.title} to watchlist.")
+
+
+@watchlist_app.command()
+def remove(pr: int):
+    """Remove PR from watchlist."""
+    wl = Watchlist.from_file()
+    if pr not in wl:
+        print(f"#{pr} not in watchlist.")
+        return
+    wl.remove(pr)
+    wl.to_file()
+    print(f"Removed #{pr} from watchlist.")
+
+
 def main():
     app()
 

nixprstatus/output.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class OutputFormat(str, Enum):
+    CONSOLE = "console"
+    JSON = "json"

nixprstatus/pr.py

@@ -1,5 +1,8 @@
 import requests
 from pydantic import BaseModel
+from rich.console import Console
+
+from nixprstatus.output import OutputFormat
 
 DEFAULT_HEADERS = {
     "Accept": "application/vnd.github.text+json",
@@ -9,9 +12,30 @@ BACKPORT_LABEL = "backport release-24.05"
 
 
 class PRStatus(BaseModel):
+    title: str
     merged: bool
     branches: dict[str, bool]
 
+    def print(self, format: OutputFormat = OutputFormat.CONSOLE):
+        match format:
+            case OutputFormat.JSON:
+                print(self.model_dump_json())
+            case OutputFormat.CONSOLE:
+                console = Console(highlight=False)
+                console.print(f"{self.title}\n")
+                merged = ":white_check_mark: merged" if self.merged else ":x: merged"
+                console.print(merged)
+
+                for branch in self.branches:
+                    output = (
+                        f":white_check_mark: {branch}"
+                        if self.branches[branch]
+                        else f":x: {branch}"
+                    )
+                    console.print(output)
+            case _:
+                raise ValueError(f"Unknown format: {format}")
+
 
 def commit_in_branch(commit_sha: str, branch: str) -> bool:
     url = f"https://api.github.com/repos/NixOS/nixpkgs/compare/{branch}...{commit_sha}"
@@ -32,6 +56,14 @@ def commits_since(first_ref: str, last_ref: str) -> int:
     return commit_response.json()["behind_by"]
 
 
+def get_pr(pr: int) -> dict:
+    url = f"https://api.github.com/repos/NixOS/nixpkgs/pulls/{pr}"
+    pr_response = requests.get(url, headers=DEFAULT_HEADERS)
+    pr_response.raise_for_status()
+
+    return pr_response.json()
+
+
 def pr_merge_status(
     pr: int, branches: list[str] = DEFAULT_BRANCHES, check_backport: bool = True
 ) -> PRStatus:
@@ -40,10 +72,13 @@ def pr_merge_status(
     pr_response.raise_for_status()
 
     pr_data = pr_response.json()
+    title = pr_data["title"]
 
     merged = pr_data["merged"]
     if merged is False:
-        return PRStatus(merged=False, branches={branch: False for branch in branches})
+        return PRStatus(
+            title=title, merged=False, branches={branch: False for branch in branches}
+        )
 
     commit_sha = pr_data.get("merge_commit_sha")
 
@@ -54,6 +89,13 @@ def pr_merge_status(
 
     results = {}
 
+    # Check if base branch is in our list, if it is
+    # no need to call commit_in_branch
+    merge_base_branch = pr_data.get("base", {}).get("ref")
+    if merge_base_branch in branches:
+        results[merge_base_branch] = True
+        branches.remove(merge_base_branch)
+
     for branch in branches:
         in_branch = commit_in_branch(commit_sha, branch)
         results[branch] = in_branch
@@ -77,11 +119,11 @@ def pr_merge_status(
                 backport_sha = backport_response.json().get("merge_commit_sha")
                 if backport_sha is None:
                     results[f"nixos-24.05 (#{backport_pr})"] = False
-                    return PRStatus(merged=True, branches=results)
+                    return PRStatus(title=title, merged=True, branches=results)
 
                 results.pop("nixos-24.05")
                 results[f"nixos-24.05 (#{backport_pr})"] = commit_in_branch(
                     backport_sha, "nixos-24.05"
                 )
 
-    return PRStatus(merged=True, branches=results)
+    return PRStatus(title=title, merged=True, branches=results)

nixprstatus/watchlist.py

@@ -0,0 +1,88 @@
+import json
+import os
+from pathlib import Path
+from pydantic import BaseModel
+from rich.console import Console
+
+from nixprstatus.pr import get_pr
+from nixprstatus.output import OutputFormat
+
+
+class PRInfo(BaseModel):
+    pr: int
+    title: str
+
+
+class Watchlist(BaseModel):
+    prs: list[PRInfo]
+
+    @classmethod
+    def from_file(cls, path: str | None = None) -> "Watchlist":
+        if not path:
+            path = _default_path()
+
+        p = Path(path).expanduser()
+
+        if not p.exists():
+            return cls(prs=[])
+
+        with open(p, "r") as f:
+            data = json.load(f)
+        return cls(**data)
+
+    def to_file(self, path: str | None = None):
+        if not path:
+            _ensure_default_path()
+            path = _default_path()
+
+        p = Path(path).expanduser()
+
+        with open(p, "w") as f:
+            f.write(self.model_dump_json())
+
+    def add_pr(self, pr: int) -> PRInfo:
+        # Lookup PR info
+        info = get_pr(pr)
+
+        title = info["title"]
+        info = PRInfo(pr=pr, title=title)
+        self.prs.append(info)
+        return info
+
+    def remove(self, pr: int):
+        self.prs = [p for p in self.prs if p.pr != pr]
+
+    def print(self, format: OutputFormat = OutputFormat.CONSOLE):
+        match format:
+            case OutputFormat.CONSOLE:
+                console = Console()
+                for pr in self.prs:
+                    console.print(f"{pr.pr}: {pr.title}")
+            case OutputFormat.JSON:
+                print(self.model_dump_json())
+            case _:
+                raise ValueError(f"Unknown format: {format}")
+
+    def pr(self, pr: int) -> PRInfo | None:
+        for p in self.prs:
+            if p.pr == pr:
+                return p
+        return None
+
+    def __contains__(self, item: PRInfo | int):
+        match item:
+            case PRInfo():
+                return any([x == item for x in self.prs])
+            case int():
+                return any([x.pr == item for x in self.prs])
+
+
+def _default_path() -> str:
+    if "XDG_STATE_HOME" in os.environ:
+        return f"{os.environ['XDG_STATE_HOME']}/nixprstatus/watchlist.json"
+    return "~/.config/nixprstatus/watchlist.json"
+
+
+def _ensure_default_path():
+    p = Path(_default_path()).expanduser()
+    p.parent.mkdir(parents=True, exist_ok=True)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nixprstatus"
-version = "0.1.1"
+version = "0.1.5"
 description = "Nixpkgs PR status checker"
 authors = ["Torjus Håkestad <torjus@usit.uio.no>"]
 license = "MIT"

tests/helpers/mocks.py

@@ -0,0 +1,32 @@
+import requests
+import json
+
+
+def mocked_requests_get(*args, **kwargs):
+    class MockedResponse:
+        def __init__(self, json_data, status_code):
+            self.json_data = json_data
+            self.status_code = status_code
+
+        def json(self):
+            return json.loads(self.json_data)
+
+        def raise_for_status(self):
+            if self.status_code not in [200, 201]:
+                raise requests.exceptions.HTTPError()
+
+    if "pulls" in args[0]:
+        pr = args[0].split("/")[-1]
+        with open(f"tests/fixtures/pulls_{pr}.json") as f:
+            data = f.read()
+        return MockedResponse(data, 200)
+    elif "compare" in args[0]:
+        branch, commit_sha = args[0].split("/")[-1].split("...")
+        with open(f"tests/fixtures/compare_{branch}_{commit_sha}.json") as f:
+            data = f.read()
+        return MockedResponse(data, 200)
+    elif "comments" in args[0]:
+        pr = args[0].split("/")[-2]
+        with open(f"tests/fixtures/comments_{pr}.json") as f:
+            data = f.read()
+        return MockedResponse(data, 200)

tests/test_pr.py

@@ -1,39 +1,8 @@
 import unittest
 import unittest.mock
-import requests
-import json
 
 from nixprstatus.pr import commit_in_branch, pr_merge_status, commits_since
-
+from tests.helpers.mocks import mocked_requests_get
-
-def mocked_requests_get(*args, **kwargs):
-    class MockedResponse:
-        def __init__(self, json_data, status_code):
-            self.json_data = json_data
-            self.status_code = status_code
-
-        def json(self):
-            return json.loads(self.json_data)
-
-        def raise_for_status(self):
-            if self.status_code not in [200, 201]:
-                raise requests.exceptions.HTTPError()
-
-    if "pulls" in args[0]:
-        pr = args[0].split("/")[-1]
-        with open(f"tests/fixtures/pulls_{pr}.json") as f:
-            data = f.read()
-        return MockedResponse(data, 200)
-    elif "compare" in args[0]:
-        branch, commit_sha = args[0].split("/")[-1].split("...")
-        with open(f"tests/fixtures/compare_{branch}_{commit_sha}.json") as f:
-            data = f.read()
-        return MockedResponse(data, 200)
-    elif "comments" in args[0]:
-        pr = args[0].split("/")[-2]
-        with open(f"tests/fixtures/comments_{pr}.json") as f:
-            data = f.read()
-        return MockedResponse(data, 200)
 
 
 class TestPRMergeStatus(unittest.TestCase):
@@ -57,6 +26,28 @@ class TestPRMergeStatus(unittest.TestCase):
         self.assertTrue(res.merged)
         self.assertTrue(res.branches["nixos-24.05 (#346022)"])
 
+    @unittest.mock.patch("requests.get", side_effect=mocked_requests_get)
+    def test_pr_merge_status_title_345769(self, mock_get):
+        pr = 345769
+        branches = ["nixos-24.05"]
+        expected_title = "Firefox: 130.0.1 -> 131.0; 128.2.0esr -> 128.3.0esr; 115.15.0esr -> 115.16.0esr"
+
+        res = pr_merge_status(pr, branches, check_backport=True)
+        self.assertEqual(res.title, expected_title)
+
+    @unittest.mock.patch("requests.get", side_effect=mocked_requests_get)
+    def test_pr_merge_status_no_check_master_345583(self, mock_get):
+        pr = 345583
+        branches = ["master", "nixos-unstable", "nixos-24.05"]
+        master_compare_url = "https://api.github.com/repos/NixOS/nixpkgs/compare/master...2c5fac3edf2d00d948253e392ec1604b29b38f14"
+
+        res = pr_merge_status(pr, branches, check_backport=False)
+        self.assertTrue(res.merged)
+        self.assertTrue(res.branches["master"])
+
+        urls_called = [call[0][0] for call in mock_get.call_args_list]
+        self.assertFalse(master_compare_url in urls_called)
+
 
 class TestCommitInBranch(unittest.TestCase):
     @unittest.mock.patch("requests.get", side_effect=mocked_requests_get)

tests/test_watchlist.py

@@ -0,0 +1,40 @@
+from nixprstatus.watchlist import Watchlist, PRInfo
+from tempfile import TemporaryDirectory
+import unittest
+
+from tests.helpers.mocks import mocked_requests_get
+
+
+class TestWatchlist(unittest.TestCase):
+    def test_save_load(self):
+        with TemporaryDirectory() as d:
+            filename = f"{d}/test.json"
+
+            watchlist = Watchlist(prs=[PRInfo(pr=1, title="PR 1")])
+            watchlist.to_file(filename)
+
+            # Check that the file was written correctly
+            with open(filename, "r") as f:
+                self.assertEqual(watchlist.model_dump_json(), f.read())
+
+            # Check that the file can be read back
+            loaded = Watchlist.from_file(filename)
+            self.assertEqual(watchlist, loaded)
+
+    @unittest.mock.patch("requests.get", side_effect=mocked_requests_get)
+    def test_add_pr(self, mock_get):
+        w = Watchlist(prs=[])
+        w.add_pr(345583)
+        self.assertEqual(len(w.prs), 1)
+        self.assertEqual(w.prs[0].title, "wireshark: 4.2.6 -> 4.2.7")
+
+    def test_get_pr(self):
+        w = Watchlist(prs=[PRInfo(pr=1, title="PR 1")])
+        self.assertEqual(w.pr(1), PRInfo(pr=1, title="PR 1"))
+        self.assertEqual(w.pr(2), None)
+
+    def test_contains(self):
+        w = Watchlist(prs=[PRInfo(pr=1, title="PR 1")])
+        self.assertIn(PRInfo(pr=1, title="PR 1"), w)
+        self.assertIn(1, w)
+        self.assertNotIn(2, w)