Skip to content

Instantly share code, notes, and snippets.

@dlmanning
Last active April 6, 2023 10:12
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save dlmanning/fc28cf8a022bd8ae15760f6debe8e2c4 to your computer and use it in GitHub Desktop.
Save dlmanning/fc28cf8a022bd8ae15760f6debe8e2c4 to your computer and use it in GitHub Desktop.
Little utility to list all PR's waiting for your review in a specified repo
[package]
name = "gh-pr-cli"
version = "0.1.0"
edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
async-openai = "0.10.2"
dotenvy = "0.15.7"
octocrab = "0.19.0"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0.93"
tracing = "0.1.37"
tracing-subscriber = { version = "0.3.16", features = ["env-filter"] }
tokio = { version = "1.27.0", features = ["full"] }
url = "2.3.1"
chrono = "0.4.24"
colored = "2.0.0"
prettytable-rs = "0.10.0"
clap = { version = "4.2.1", features = ["derive"] }
futures = "0.3.28"
#[macro_use]
extern crate prettytable;
use std::{
collections::{HashMap, HashSet},
hash::{Hash, Hasher},
};
use clap::Parser;
use colored::*;
use dotenvy::dotenv;
use futures::{stream::FuturesUnordered, StreamExt};
use octocrab::{
models::{pulls::Comment, teams::Team, User},
params::State,
Octocrab, Page, Result,
};
use tokio::task::JoinHandle;
#[derive(Parser)]
#[command(author, version, about)]
struct Cli {
#[clap(short = 'r', long = "repo")]
repo: String,
#[clap(short = 'c', long = "comments", default_value = "false")]
comments: bool,
}
type GhApiPullRequest = octocrab::models::pulls::PullRequest;
struct PullRequest<'a>(&'a GhApiPullRequest);
impl Hash for PullRequest<'_> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.0.number.hash(state);
}
}
impl PartialEq for PullRequest<'_> {
fn eq(&self, other: &Self) -> bool {
self.0.number == other.0.number
}
}
impl Eq for PullRequest<'_> {}
#[tokio::main]
async fn main() -> Result<()> {
dotenv().ok();
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
.with_ansi(true)
.init();
let token = std::env::var("GH_TOKEN").expect("GH_TOKEN must be set");
let octocrab = octocrab::Octocrab::builder()
.personal_token(token)
.build()?;
let cli = Cli::parse();
let repo = cli.repo;
let mut repo_parts = repo.split('/');
let (owner, repo) = match (repo_parts.next(), repo_parts.next()) {
(Some(owner), Some(repo)) => (owner.to_string(), repo.to_string()),
(Some(repo), None) => (
{
let me = octocrab.current().user().await?;
me.login.to_owned()
},
repo.to_string(),
),
_ => {
eprintln!("Invalid repository name");
std::process::exit(1);
}
};
let teams = octocrab.get(format!("/user/teams"), None::<&()>).await?;
let teams: Vec<Team> = serde_json::from_value(teams).unwrap();
let my_teams: Vec<&Team> = teams
.iter()
.filter(|&team| {
team.organization
.as_ref()
.map_or(false, |org| org.login == owner)
})
.collect();
let me = octocrab.current().user().await?;
let pulls = octocrab.pulls(&owner, &repo);
println!("{} {}", "Fetching PRs for".bold(), repo.bold());
let prs = pulls.list().per_page(50).state(State::Open).send().await?;
let mut table1 = prettytable::Table::new();
table1.set_titles(row![
"PR #",
"Last Updated",
"Title",
"URL",
"Author",
"+/-"
]);
let mut prs_concerning_me: HashSet<PullRequest> = HashSet::new();
println!("Checking for PRs assigned to me or my teams...");
for pr in &prs {
let show = process_pr(&pr, my_teams.clone(), &me);
if show {
prs_concerning_me.insert(PullRequest(pr));
}
}
println!("Checking for PRs mentioning me...");
if cli.comments {
let prs_mentioning_me =
comments_mention_me(&octocrab, &prs, &owner, &repo, &me.login).await?;
prs_mentioning_me.into_iter().for_each(|pr| {
prs_concerning_me.insert(PullRequest(pr));
});
}
let pr_additions_and_deletions =
get_additions_deletions(&octocrab, &owner, &repo, &prs_concerning_me).await?;
let mut ordered_prs = prs_concerning_me.into_iter().collect::<Vec<PullRequest>>();
ordered_prs.sort_by(|a, b| {
a.0.updated_at
.cmp(&b.0.updated_at)
.reverse()
.then(a.0.number.cmp(&b.0.number))
});
for pr in ordered_prs {
let pr = pr.0;
let (additions, deletions) = pr_additions_and_deletions
.get(&pr.number)
.map_or((0, 0), |(a, d)| (*a, *d));
table1.add_row(make_table_row(pr, (additions, deletions)));
}
table1.printstd();
Ok(())
}
fn process_pr(pr: &GhApiPullRequest, my_teams: Vec<&Team>, me: &User) -> bool {
let requested_review = pr
.requested_reviewers
.iter()
.flatten()
.any(|reviewer| reviewer.login == me.login);
let mentions_me = pr
.body
.as_ref()
.unwrap_or(&"".to_string())
.contains(&format!("@{}", me.login));
let assigned_to_my_team = pr.requested_teams.iter().flatten().any(|team| {
my_teams
.iter()
.any(|my_team| team.id.map_or(false, |id| id == my_team.id))
});
let assigned_to_me = pr
.assignees
.iter()
.flatten()
.any(|assignee| assignee.login == me.login);
requested_review || mentions_me || assigned_to_my_team || assigned_to_me
}
async fn comments_mention_me<'a>(
octocrab: &Octocrab,
prs: &'a Page<GhApiPullRequest>,
owner: &'a String,
repo: &'a String,
login: &'a String,
) -> Result<Vec<&'a GhApiPullRequest>> {
let mut result: Vec<&GhApiPullRequest> = Vec::new();
let mut handles: FuturesUnordered<JoinHandle<Result<(u64, Page<Comment>)>>> =
FuturesUnordered::new();
for pr in prs.clone() {
let octocrab = octocrab.clone();
let owner = owner.clone();
let repo = repo.clone();
let handle = tokio::spawn(async move {
let comments = octocrab
.pulls(owner, repo)
.list_comments(Some(pr.number))
.send()
.await?;
Ok((pr.number, comments))
});
handles.push(handle);
}
while let Some(res) = handles.next().await {
let (pr, comments) = res.unwrap()?;
if comments.items.iter().any(|comment| {
comment.body.contains(&format!("@{}", login))
|| if let Some(user) = &comment.user {
user.login.eq(login)
} else {
false
}
}) {
let pr = prs.into_iter().find(|&p| p.number == pr).unwrap();
result.push(&pr);
}
}
Ok(result)
}
async fn get_additions_deletions(
octocrab: &Octocrab,
owner: &String,
repo: &String,
prs: &HashSet<PullRequest<'_>>,
) -> Result<HashMap<u64, (u64, u64)>> {
let mut results: HashMap<u64, (u64, u64)> = HashMap::new();
let mut handles: FuturesUnordered<JoinHandle<Result<(u64, (u64, u64))>>> =
FuturesUnordered::new();
let pr_numbers: Vec<u64> = prs.iter().map(|pr| pr.0.number).collect();
for pr_number in pr_numbers {
let octocrab = octocrab.clone();
let owner = owner.clone();
let repo = repo.clone();
let handle = tokio::spawn(async move {
let files = octocrab.pulls(owner, repo).list_files(pr_number).await?;
let (additions, deletions) = files.into_iter().fold((0, 0), |(a, d), file| {
(a + file.additions, d + file.deletions)
});
Ok((pr_number, (additions, deletions)))
});
handles.push(handle);
}
while let Some(res) = handles.next().await {
let (pr, (a, d)) = res.unwrap()?;
results.insert(pr, (a, d));
}
Ok(results)
}
fn make_table_row(
pr: &octocrab::models::pulls::PullRequest,
(additions, deletions): (u64, u64),
) -> prettytable::Row {
row![
pr.number.to_string().bold(),
pr.updated_at.map_or("Unknown".to_string(), |date| {
date.with_timezone(&chrono::Local)
.format("%Y-%m-%d %H:%M %Z")
.to_string()
.yellow()
.to_string()
}),
pr.title.as_ref().unwrap_or(&"No title".to_string()).bold(),
pr.html_url
.as_ref()
.map_or("Unknown".to_string(), |url| url.to_string())
.underline(),
pr.user
.as_ref()
.map_or("Unknown".to_string(), |u| format!("@{}", u.login))
.green(),
format!(
"+{}/-{}",
additions.to_string().green(),
deletions.to_string().red()
)
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment