use anyhow::{bail, Context, Result};
use clap::Parser;
use nix_cache_pin_lib::wheel_pins::{
diff, latest_available, read_sidecar, resolve, update, verify, UpdateStatus, WheelPinConfig,
};
use std::path::PathBuf;
#[derive(Debug, Parser)]
#[command(
about = "Pin prebuilt Python wheels (e.g. PyTorch ROCm/XPU) to exact files and sha256 hashes"
)]
struct Cli {
#[arg(long)]
config: PathBuf,
#[arg(long)]
output_file: Option<PathBuf>,
#[arg(short = 'n', long)]
dry_run: bool,
#[arg(long, conflicts_with_all = ["dry_run", "verify_sidecar", "plan"])]
check_current: bool,
#[arg(long, conflicts_with_all = ["dry_run", "verify_sidecar"])]
plan: bool,
#[arg(long, value_name = "SIDECAR")]
verify_sidecar: Option<PathBuf>,
#[arg(long)]
json: bool,
}
#[tokio::main]
async fn main() -> Result<()> {
let cli = Cli::parse();
if cli.json && !cli.check_current {
bail!("--json is only valid with --check-current");
}
let raw =
std::fs::read(&cli.config).with_context(|| format!("read {}", cli.config.display()))?;
let mut config: WheelPinConfig =
serde_json::from_slice(&raw).with_context(|| format!("parse {}", cli.config.display()))?;
if let Some(output) = cli.output_file {
config.output_file = output;
}
if let Some(sidecar_path) = cli.verify_sidecar {
let sidecar = read_sidecar(&sidecar_path)?.with_context(|| {
format!(
"{} is missing; run 'nix run .#cache-pin-wheel-pins-{}'",
sidecar_path.display(),
config.name
)
})?;
verify(&config, &sidecar)?;
println!(
"wheel pin {}: {} wheel(s) verified",
config.name,
sidecar.wheels.len()
);
return Ok(());
}
let client = nix_cache_pin_lib::http::client()?;
if cli.plan {
let current = read_sidecar(&config.output_file)?;
let resolved = resolve(&config, client).await?;
verify(&config, &resolved)?;
let plan = serde_json::json!({
"schemaVersion": 1,
"name": config.name,
"changes": diff(current.as_ref(), &resolved),
"resolved": resolved,
});
println!("{}", serde_json::to_string_pretty(&plan)?);
return Ok(());
}
if cli.check_current {
let current = read_sidecar(&config.output_file)?;
let resolved = resolve(&config, client.clone()).await?;
let latest = latest_available(&config, client).await?;
let pinned = |p: &str| {
current.as_ref().and_then(|s| {
s.wheels
.iter()
.find(|w| w.pname == p)
.map(|w| w.version.clone())
})
};
let allowed = |p: &str| {
resolved
.wheels
.iter()
.find(|w| w.pname == p)
.map(|w| w.version.clone())
};
let rows: Vec<_> = latest
.iter()
.map(|(p, newest)| {
serde_json::json!({
"pname": p, "pinned": pinned(p), "allowed": allowed(p), "newest": newest,
})
})
.collect();
let changes = diff(current.as_ref(), &resolved);
if cli.json {
println!(
"{}",
serde_json::to_string_pretty(&serde_json::json!({
"schemaVersion": 1, "name": config.name, "packages": rows, "pendingChanges": changes,
}))?
);
} else {
println!("=== cache-pin wheel-pins: {} ===", config.name);
for row in &rows {
let p = row["pname"].as_str().unwrap_or_default();
let fmt = |v: &serde_json::Value| v.as_str().unwrap_or("-").to_owned();
let blocked = row["allowed"] != row["newest"];
println!(
" {p}: pinned {} | constraints allow {} | newest {}{}",
fmt(&row["pinned"]),
fmt(&row["allowed"]),
fmt(&row["newest"]),
if blocked {
" (blocked by constraint)"
} else {
""
}
);
}
if changes.is_empty() {
println!(" Sidecar is current");
} else {
println!(
" {} pinned wheel(s) would change; run without --check-current to update",
changes.len()
);
}
}
return Ok(());
}
println!("=== cache-pin wheel-pins: {} ===", config.name);
println!(" Output: {}", config.output_file.display());
match update(&config, client, cli.dry_run).await? {
UpdateStatus::Current => println!(" Sidecar is current"),
UpdateStatus::WouldUpdate(changes) | UpdateStatus::Updated(changes)
if changes.is_empty() =>
{
println!(" Sidecar metadata refreshed (no version changes)")
}
UpdateStatus::WouldUpdate(changes) => {
for c in &changes {
println!(" would change {}: {:?} -> {:?}", c.pname, c.from, c.to);
}
}
UpdateStatus::Updated(changes) => {
for c in &changes {
println!(" {}: {:?} -> {:?}", c.pname, c.from, c.to);
}
}
}
Ok(())
}