Skip to content

Commit

Permalink
feat: Improve UX by adding console output
Browse files Browse the repository at this point in the history
  • Loading branch information
delsner committed Jun 7, 2024
1 parent d4bef2c commit 3b219d7
Show file tree
Hide file tree
Showing 4 changed files with 148 additions and 16 deletions.
2 changes: 2 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
mod pack;
mod unpack;
mod util;

pub use pack::{pack, PackOptions};
use rattler_conda_types::Platform;
use serde::{Deserialize, Serialize};
pub use unpack::{unarchive, unpack, UnpackOptions};
pub use util::{create_progress_bar, get_size, InstallationProgressReporter};

pub const CHANNEL_DIRECTORY_NAME: &str = "channel";
pub const PIXI_PACK_METADATA_PATH: &str = "pixi-pack.json";
Expand Down
42 changes: 29 additions & 13 deletions src/pack.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ use std::{
};

use fxhash::FxHashMap;
use indicatif::HumanBytes;
use rattler_index::{package_record_from_conda, package_record_from_tar_bz2};
use tokio::{
fs::{self, create_dir_all, File},
Expand All @@ -13,14 +14,16 @@ use tokio::{

use anyhow::Result;
use futures::{stream, StreamExt, TryFutureExt, TryStreamExt};
use indicatif::ProgressStyle;
use rattler_conda_types::{package::ArchiveType, ChannelInfo, PackageRecord, Platform, RepoData};
use rattler_lock::{CondaPackage, LockFile, Package};
use rattler_networking::{AuthenticationMiddleware, AuthenticationStorage};
use reqwest_middleware::ClientWithMiddleware;
use tokio_tar::Builder;

use crate::{PixiPackMetadata, CHANNEL_DIRECTORY_NAME, PIXI_PACK_METADATA_PATH};
use crate::{
create_progress_bar, get_size, PixiPackMetadata, CHANNEL_DIRECTORY_NAME,
PIXI_PACK_METADATA_PATH,
};
use anyhow::anyhow;

/// Options for packing a pixi environment.
Expand Down Expand Up @@ -89,31 +92,26 @@ pub async fn pack(options: PackOptions) -> Result<()> {

// Download packages to temporary directory.
tracing::info!(
"Downloading {} packages",
"Downloading {} packages...",
conda_packages_from_lockfile.len()
);
let bar = indicatif::ProgressBar::new(conda_packages_from_lockfile.len() as u64);
bar.set_style(
ProgressStyle::with_template(
"[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}",
)
.expect("could not set progress style")
.progress_chars("##-"),
println!(
"⏳ Downloading {} packages...",
conda_packages_from_lockfile.len()
);
let bar = create_progress_bar(conda_packages_from_lockfile.len() as u64);

stream::iter(conda_packages_from_lockfile.iter())
.map(Ok)
.try_for_each_concurrent(50, |package| async {
download_package(&client, package, &channel_dir).await?;

bar.inc(1);

Ok(())
})
.await
.map_err(|e: anyhow::Error| anyhow!("could not download package: {}", e))?;

bar.finish();
bar.finish_and_clear();

let mut conda_packages: Vec<(String, PackageRecord)> = Vec::new();

Expand All @@ -133,6 +131,8 @@ pub async fn pack(options: PackOptions) -> Result<()> {
.map(|(p, t)| (PathBuf::from(format!("{}{}", p, t.extension())), t))
})
.collect();

tracing::info!("Injecting {} packages", injected_packages.len());
for (path, archive_type) in injected_packages {
// step 1: Derive PackageRecord from index.json inside the package
let package_record = match archive_type {
Expand All @@ -157,23 +157,39 @@ pub async fn pack(options: PackOptions) -> Result<()> {
}

// Create `repodata.json` files.
tracing::info!("Creating repodata.json files");
create_repodata_files(conda_packages.iter(), &channel_dir).await?;

// Add pixi-pack.json containing metadata.
tracing::info!("Creating pixi-pack.json file");
let metadata_path = output_folder.path().join(PIXI_PACK_METADATA_PATH);
let mut metadata_file = File::create(&metadata_path).await?;

let metadata = serde_json::to_string_pretty(&options.metadata)?;
metadata_file.write_all(metadata.as_bytes()).await?;

// Create environment file.
tracing::info!("Creating environment.yml file");
create_environment_file(output_folder.path(), conda_packages.iter().map(|(_, p)| p)).await?;

// Pack = archive the contents.
tracing::info!("Creating archive at {}", options.output_file.display());
archive_directory(output_folder.path(), &options.output_file)
.await
.map_err(|e| anyhow!("could not archive directory: {}", e))?;

let output_size = HumanBytes(get_size(&options.output_file)?).to_string();
tracing::info!(
"Created pack at {} with size {}.",
options.output_file.display(),
output_size
);
println!(
"📦 Created pack at {} with size {}.",
options.output_file.display(),
output_size
);

Ok(())
}

Expand Down
23 changes: 20 additions & 3 deletions src/unpack.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,8 @@ use tokio_tar::Archive;
use url::Url;

use crate::{
PixiPackMetadata, CHANNEL_DIRECTORY_NAME, DEFAULT_PIXI_PACK_VERSION, PIXI_PACK_METADATA_PATH,
InstallationProgressReporter, PixiPackMetadata, CHANNEL_DIRECTORY_NAME,
DEFAULT_PIXI_PACK_VERSION, PIXI_PACK_METADATA_PATH,
};

/// Options for unpacking a pixi environment.
Expand All @@ -41,6 +42,7 @@ pub async fn unpack(options: UnpackOptions) -> Result<()> {

let channel_directory = unpack_dir.join(CHANNEL_DIRECTORY_NAME);

tracing::info!("Unarchiving pack to {}", unpack_dir.display());
unarchive(&options.pack_file, &unpack_dir)
.await
.map_err(|e| anyhow!("Could not unarchive: {}", e))?;
Expand All @@ -49,10 +51,12 @@ pub async fn unpack(options: UnpackOptions) -> Result<()> {

let target_prefix = options.output_directory.join("env");

tracing::info!("Creating prefix at {}", target_prefix.display());
create_prefix(&channel_directory, &target_prefix)
.await
.map_err(|e| anyhow!("Could not create prefix: {}", e))?;

tracing::info!("Generating activation script");
create_activation_script(
&options.output_directory,
&target_prefix,
Expand All @@ -61,6 +65,15 @@ pub async fn unpack(options: UnpackOptions) -> Result<()> {
.await
.map_err(|e| anyhow!("Could not create activation script: {}", e))?;

tracing::info!(
"Finished unpacking to {}.",
options.output_directory.display(),
);
println!(
"💫 Finished unpacking to {}.",
options.output_directory.display()
);

Ok(())
}

Expand Down Expand Up @@ -158,8 +171,6 @@ async fn create_prefix(channel_dir: &Path, target_prefix: &Path) -> Result<()> {
// extract packages to cache
let package_cache = PackageCache::new(cache_dir);

let installer = Installer::default();

let repodata_records: Vec<RepoDataRecord> = stream::iter(packages)
.map(|(file_name, package_record)| {
let cache_key = CacheKey::from(&package_record);
Expand Down Expand Up @@ -198,6 +209,12 @@ async fn create_prefix(channel_dir: &Path, target_prefix: &Path) -> Result<()> {
.await?;

// Invariant: all packages are in the cache
tracing::info!("Installing {} packages", repodata_records.len());
println!("⏳ Installing {} packages...", repodata_records.len());

let installer = Installer::default().with_reporter(InstallationProgressReporter::new(
repodata_records.len() as u64,
));
installer
.with_package_cache(package_cache)
.install(&target_prefix, repodata_records)
Expand Down
97 changes: 97 additions & 0 deletions src/util.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
use std::path::Path;

use indicatif::{ProgressBar, ProgressStyle};
use rattler::install::Reporter;
use rattler_conda_types::RepoDataRecord;

/// Create a progress bar with default style.
pub fn create_progress_bar(length: u64) -> ProgressBar {
ProgressBar::new(length).with_style(
ProgressStyle::with_template(
"[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}",
)
.expect("could not set progress style")
.progress_chars("##-"),
)
}

/// Progress reporter for installing packages.
pub struct InstallationProgressReporter {
pb: ProgressBar,
}

impl InstallationProgressReporter {
pub fn new(length: u64) -> Self {
Self {
pb: create_progress_bar(length),
}
}
}

impl Reporter for InstallationProgressReporter {
fn on_transaction_start(
&self,
_transaction: &rattler::install::Transaction<
rattler_conda_types::PrefixRecord,
RepoDataRecord,
>,
) {
}

fn on_transaction_operation_start(&self, _operation: usize) {}
fn on_download_start(&self, cache_entry: usize) -> usize {
cache_entry
}

fn on_download_completed(&self, _download_idx: usize) {}

fn on_link_start(&self, operation: usize, _record: &RepoDataRecord) -> usize {
operation
}

fn on_link_complete(&self, _index: usize) {}

fn on_transaction_operation_complete(&self, _operation: usize) {
self.pb.inc(1);
}

fn on_populate_cache_start(&self, operation: usize, _record: &RepoDataRecord) -> usize {
operation
}

fn on_validate_start(&self, cache_entry: usize) -> usize {
cache_entry
}

fn on_validate_complete(&self, _validate_idx: usize) {}

fn on_download_progress(&self, _download_idx: usize, _progress: u64, _total: Option<u64>) {}

fn on_populate_cache_complete(&self, _cache_entry: usize) {}

fn on_unlink_start(
&self,
operation: usize,
_record: &rattler_conda_types::PrefixRecord,
) -> usize {
operation
}

fn on_unlink_complete(&self, _index: usize) {}

fn on_transaction_complete(&self) {
self.pb.finish_and_clear();
}
}

/// Get the size of a file or directory in bytes.
pub fn get_size<P: AsRef<Path>>(path: P) -> std::io::Result<u64> {
let metadata = std::fs::metadata(&path)?;
let mut size = metadata.len();
if metadata.is_dir() {
for entry in std::fs::read_dir(&path)? {
size += get_size(entry?.path())?;
}
}
Ok(size)
}

0 comments on commit 3b219d7

Please sign in to comment.