Skip to content

Commit

Permalink
finalize bar on stderr with newline
Browse files Browse the repository at this point in the history
  • Loading branch information
robfitzgerald committed Dec 20, 2024
1 parent 09a9d36 commit 84c824d
Showing 1 changed file with 46 additions and 36 deletions.
82 changes: 46 additions & 36 deletions rust/routee-compass-core/src/util/fs/read_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,41 +16,6 @@ type RawCallback<'a> = Option<Box<dyn FnMut() + 'a>>;
/// if COMPASS_PROGRESS=false, the bar is deactivated, otherwise it runs.
const COMPASS_PROGRESS: &str = "COMPASS_PROGRESS";

/// reads from a CSV into an iterator of T records.
/// building the iterator may fail with an io::Error.
/// each row hasn't yet been decoded so it is provided in a Result<T, csv::Error>
///
pub fn iterator_from_csv<'a, F, T>(
filepath: F,
has_headers: bool,
mut row_callback: CsvCallback<'a, T>,
) -> Result<Box<dyn Iterator<Item = Result<T, csv::Error>> + 'a>, io::Error>
where
F: AsRef<Path>,
T: serde::de::DeserializeOwned + 'a,
{
let f = File::open(filepath.as_ref())?;
let r: Box<dyn io::Read> = if fs_utils::is_gzip(filepath) {
Box::new(BufReader::new(GzDecoder::new(f)))
} else {
Box::new(f)
};
let reader = ReaderBuilder::new()
.has_headers(has_headers)
.trim(csv::Trim::Fields)
.from_reader(r)
.into_deserialize::<T>()
.inspect(move |r| {
if let Ok(t) = r {
if let Some(cb) = &mut row_callback {
cb(t);
}
}
});

Ok(Box::new(reader))
}

/// reads a csv file into a vector. not space-optimized since size is not
/// known.
pub fn from_csv<'a, T>(
Expand All @@ -63,6 +28,7 @@ where
T: serde::de::DeserializeOwned + 'a,
{
let bar_opt = build_bar(progress);
let finalize_bar = bar_opt.is_some();

let row_callback: CsvCallback<'a, T> = match (callback, bar_opt) {
(None, None) => None,
Expand All @@ -82,6 +48,11 @@ where
.into_iter()
.collect::<Result<Vec<T>, csv::Error>>()?
.into_boxed_slice();

if finalize_bar {
eprintln!();
}

Ok(result)
}

Expand All @@ -100,6 +71,7 @@ where
F: AsRef<Path>,
{
let bar_opt = build_bar(progress);
let finalize_bar = bar_opt.is_some();

let row_callback: RawCallback = match (callback, bar_opt) {
(None, None) => None,
Expand All @@ -113,11 +85,49 @@ where
})),
};

if fs_utils::is_gzip(filepath.as_ref()) {
let result = if fs_utils::is_gzip(filepath.as_ref()) {
Ok(read_gzip(filepath, op, row_callback)?)
} else {
Ok(read_regular(filepath, op, row_callback)?)
};
if finalize_bar {
eprintln!();
}
result
}

/// reads from a CSV into an iterator of T records.
/// building the iterator may fail with an io::Error.
/// each row hasn't yet been decoded so it is provided in a Result<T, csv::Error>
pub fn iterator_from_csv<'a, F, T>(
filepath: F,
has_headers: bool,
mut row_callback: CsvCallback<'a, T>,
) -> Result<Box<dyn Iterator<Item = Result<T, csv::Error>> + 'a>, io::Error>
where
F: AsRef<Path>,
T: serde::de::DeserializeOwned + 'a,
{
let f = File::open(filepath.as_ref())?;
let r: Box<dyn io::Read> = if fs_utils::is_gzip(filepath) {
Box::new(BufReader::new(GzDecoder::new(f)))
} else {
Box::new(f)
};
let reader = ReaderBuilder::new()
.has_headers(has_headers)
.trim(csv::Trim::Fields)
.from_reader(r)
.into_deserialize::<T>()
.inspect(move |r| {
if let Ok(t) = r {
if let Some(cb) = &mut row_callback {
cb(t);
}
}
});

Ok(Box::new(reader))
}

fn read_regular<'a, F, T>(
Expand Down

0 comments on commit 84c824d

Please sign in to comment.