Skip to content

Commit

Permalink
Merge pull request #32 from ktwrd/feature/extract-progress
Browse files Browse the repository at this point in the history
Progress Bar for file extraction
  • Loading branch information
ktwrd authored Aug 27, 2024
2 parents 1274ed5 + 79a1701 commit 3e8b84b
Show file tree
Hide file tree
Showing 4 changed files with 199 additions and 33 deletions.
43 changes: 20 additions & 23 deletions src/ctx.rs
Original file line number Diff line number Diff line change
Expand Up @@ -281,13 +281,25 @@ impl RunnerContext
{
let tar_tmp_location = helper::get_tmp_file("data.tar".to_string());

let zstd_file = std::fs::File::open(&zstd_location)?;
let mut tar_tmp_file = std::fs::File::create_new(&tar_tmp_location)?;
zstd::stream::copy_decode(zstd_file, &tar_tmp_file)?;
tar_tmp_file = std::fs::File::open(&tar_tmp_location)?; // we do this again to make sure that the tar is properly opened.

let mut archive = tar::Archive::new(&tar_tmp_file);
let x = archive.unpack(&out_dir);
if let Err(e) =
crate::extract::decompress_zstd(zstd_location.clone(), tar_tmp_location.clone(), true)
{
debug!("{:#?}", e);
error!(
"[RunnerContext::extract_package] Failed to decompress file {} ({:})",
zstd_location, e
);
return Err(e);
}
if let Err(e) = crate::extract::unpack_tarball(tar_tmp_location.clone(), out_dir, true)
{
debug!("{:#?}", e);
error!(
"[RunnerContext::extract_package] Failed to unpack tarball {} ({:})",
tar_tmp_location, e
);
return Err(e);
}
if helper::file_exists(tar_tmp_location.clone())
{
if let Err(e) = std::fs::remove_file(tar_tmp_location.clone())
Expand All @@ -303,22 +315,7 @@ impl RunnerContext
);
}
}
match x
{
Err(e) =>
{
let xe = BeansError::TarExtractFailure {
src_file: tar_tmp_location,
target_dir: out_dir,
error: e,
backtrace: Backtrace::capture()
};
trace!("[RunnerContext::extract_package] {:}\n{:#?}", xe, xe);
sentry::capture_error(&xe);
Err(xe)
}
Ok(_) => Ok(())
}
Ok(())
}

#[cfg(target_os = "linux")]
Expand Down
9 changes: 9 additions & 0 deletions src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,15 @@ pub enum BeansError
error: std::io::Error,
backtrace: Backtrace
},
#[error("Failed to extract item {link_name} to directory {target_dir} ({error:})")]
TarUnpackItemFailure
{
src_file: String,
target_dir: String,
link_name: String,
error: std::io::Error,
backtrace: Backtrace
},
#[error("Failed to send request ({error:})")]
Reqwest
{
Expand Down
155 changes: 155 additions & 0 deletions src/extract.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,155 @@
use std::{backtrace::Backtrace,
fs::File,
io::Read};

use indicatif::{ProgressBar,
ProgressStyle};
use log::info;
use zstd::stream::read::Decoder as ZstdDecoder;

use crate::BeansError;

pub fn unpack_tarball(
tarball_location: String,
output_directory: String,
show_progress: bool
) -> Result<(), BeansError>
{
let tarball = match File::open(&tarball_location)
{
Ok(x) => x,
Err(e) =>
{
return Err(BeansError::TarExtractFailure {
src_file: tarball_location,
target_dir: output_directory,
error: e,
backtrace: Backtrace::capture()
});
}
};
let mut archive = tar::Archive::new(&tarball);
if show_progress
{
let archive_entries = match archive.entries()
{
Ok(v) => v,
Err(e) =>
{
return Err(BeansError::TarExtractFailure {
src_file: tarball_location,
target_dir: output_directory,
error: e,
backtrace: Backtrace::capture()
});
}
};
let archive_entry_count = (&archive_entries.count()).clone() as u64;
info!("Extracting {} files", archive_entry_count);

let pb = ProgressBar::new(archive_entry_count);
pb.set_style(ProgressStyle::with_template("{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos}/{len} ({eta})")
.unwrap()
.with_key("eta", |state: &indicatif::ProgressState, w: &mut dyn std::fmt::Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap())
.progress_chars("#>-"));
pb.set_message("Extracting files");

archive = tar::Archive::new(&tarball);
match archive.entries()
{
Ok(etrs) =>
{
for entry in etrs
{
match entry
{
Ok(mut x) =>
{
let ln = x.link_name();
pb.set_message("Extracting files");
let mut filename = String::new();
if let Ok(n) = ln
{
if let Some(p) = n
{
if let Some(s) = p.to_str()
{
pb.set_message(format!("{:}", s));
filename = String::from(s);
}
}
}
if let Err(e) = x.unpack_in(&output_directory)
{
return Err(BeansError::TarUnpackItemFailure {
src_file: tarball_location,
target_dir: output_directory,
link_name: filename,
error: e,
backtrace: Backtrace::capture()
});
}
pb.inc(1);
}
Err(e) =>
{
return Err(BeansError::TarExtractFailure {
src_file: tarball_location,
target_dir: output_directory,
error: e,
backtrace: Backtrace::capture()
});
}
};
}
}
Err(e) =>
{
return Err(BeansError::TarExtractFailure {
src_file: tarball_location,
target_dir: output_directory,
error: e,
backtrace: Backtrace::capture()
});
}
};
pb.finish();
}
else
{
archive.unpack(output_directory);
}
return Ok(());
}
pub fn decompress_zstd(
zstd_location: String,
output_file: String,
show_progress: bool
) -> Result<(), BeansError>
{
let zstd_file = File::open(&zstd_location)?;
let zstd_file_length = &zstd_file.metadata()?.len();
let mut tar_tmp_file = File::create_new(&output_file)?;
if show_progress
{
let decoder = ZstdDecoder::new(zstd_file)?;
// estimate extracted size as x2 since idk how to get the decompressed size with
// zstd
let pb_decompress = ProgressBar::new((zstd_file_length.clone() * 2) as u64);
pb_decompress
.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})")
.unwrap()
.with_key("eta", |state: &indicatif::ProgressState, w: &mut dyn std::fmt::Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap())
.progress_chars("#>-"));

std::io::copy(&mut pb_decompress.wrap_read(decoder), &mut tar_tmp_file)
.expect("Failed to decompress file");
pb_decompress.finish();
}
else
{
zstd::stream::copy_decode(zstd_file, &tar_tmp_file)?;
}

Ok(())
}
25 changes: 15 additions & 10 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ pub use error::*;

pub mod appvar;
pub mod butler;
pub mod extract;
pub mod flags;
pub mod gui;
pub mod logger;
Expand Down Expand Up @@ -59,33 +60,36 @@ pub fn data_dir() -> String

/// Temporary directory which is specified by `ADASTRAL_TMPDIR`.
///
/// Will return `None` when the environment variable couldn't be found, or it's an empty string.
/// Will return `None` when the environment variable couldn't be found, or it's
/// an empty string.
pub fn env_custom_tmpdir() -> Option<String>
{
let s = helper::try_get_env_var(String::from("ADASTRAL_TMPDIR"));
match s {
Some(x) => match x.trim().is_empty() {
match s
{
Some(x) => match x.trim().is_empty()
{
true => None,
false => Some(x)
},
None => s
}
}
/// Return `true` when the environment variable `BEANS_DEBUG` or `ADASTRAL_DEBUG` exists and
/// equals `1` or `true`.
/// Return `true` when the environment variable `BEANS_DEBUG` or
/// `ADASTRAL_DEBUG` exists and equals `1` or `true`.
pub fn env_debug() -> bool
{
check_env_bool("BEANS_DEBUG") || check_env_bool("ADASTRAL_DEBUG")
}
/// Return `true` when the environment variable `BEANS_HEADLESS` or `ADASTRAL_HEADLESS` exists and
/// equals `1` or `true`.
/// Return `true` when the environment variable `BEANS_HEADLESS` or
/// `ADASTRAL_HEADLESS` exists and equals `1` or `true`.
pub fn env_headless() -> bool
{
check_env_bool("BEANS_HEADLESS") || check_env_bool("ADASTRAL_HEADLESS")
}

/// Return `true` when the environment variable exists, and it's value equals `1` or `true (when
/// trimmed and made lowercase).
/// Return `true` when the environment variable exists, and it's value equals
/// `1` or `true (when trimmed and made lowercase).
fn check_env_bool<K: AsRef<std::ffi::OsStr>>(key: K) -> bool
{
std::env::var(key).is_ok_and(|x| {
Expand All @@ -111,7 +115,8 @@ pub fn has_gui_support() -> bool
}
}

if env_headless() {
if env_headless()
{
return true;
}

Expand Down

0 comments on commit 3e8b84b

Please sign in to comment.