use std::{
fmt::Display,
fs,
io::{BufWriter, Write},
path::Path,
};
use serde::Serialize;
use thiserror::Error;
use crate::{
extra_keys::partition_extra_keys,
files::{read_nb, read_nb_stdin, NBReadError, NBWriteError},
schema::RawNotebook,
settings::Settings,
utils::{get_value_child, pop_cell_key, pop_meta_key},
};
use serde_json::Value;
pub fn strip_nb(mut nb: RawNotebook, settings: &Settings) -> (RawNotebook, bool) {
let (cell_keys, meta_keys) = partition_extra_keys(&settings.extra_keys);
let nb_keep_output = get_value_child(&nb.metadata, &["keep_output"])
.and_then(Value::as_bool)
.unwrap_or(false);
let drop_output = settings.drop_output && !nb_keep_output;
let mut stripped = false;
for meta_key in meta_keys {
stripped |= pop_meta_key(&mut nb, meta_key).is_some();
}
let drop_cells: Vec<_> = nb
.cells
.iter()
.map(|c| c.should_drop(settings.drop_empty_cells, &settings.drop_tagged_cells))
.collect();
if drop_cells.iter().any(|b| *b) {
stripped = true;
let mut retained_cells = vec![];
for (cell, to_drop) in nb.cells.into_iter().zip(drop_cells.iter()) {
if !to_drop {
retained_cells.push(cell);
}
}
nb.cells = retained_cells;
}
for (i, cell) in nb.cells.iter_mut().enumerate() {
if let Some(codecell) = cell.as_codecell_mut() {
if codecell.should_clear_output(drop_output, settings.strip_init_cell)
&& !codecell.is_clear_outputs()
{
stripped = true;
codecell.clear_outputs();
}
if settings.drop_count && !codecell.is_clear_exec_count() {
stripped = true;
codecell.clear_counts();
}
}
if settings.drop_id && !cell.is_clear_id(i) {
stripped = true;
cell.set_id(Some(format!("{i}")));
}
for cell_key in &cell_keys {
stripped |= pop_cell_key(cell, cell_key).is_some();
}
}
(nb, stripped)
}
pub fn strip_single(
nb_path: &Path,
textconv: bool,
settings: &Settings,
) -> Result<StripSuccess, StripError> {
let (nb, to_stdout) = match nb_path.to_str() {
Some("-") => (read_nb_stdin()?, true),
_ => (read_nb(nb_path)?, textconv),
};
let (strip_nb, stripped) = strip_nb(nb, settings);
match (to_stdout, stripped) {
(true, _) => {
let stdout = std::io::stdout();
match write_nb(stdout, &strip_nb) {
Ok(()) => Ok(StripSuccess::from_stripped(stripped)),
#[cfg(not(tarpaulin_include))]
Err(e) => Err(e.into()),
}
}
(false, false) => Ok(StripSuccess::NoChange),
(false, true) => {
let f = fs::File::create(nb_path).map_err(NBWriteError::from)?;
let writer = BufWriter::new(f);
match write_nb(writer, &strip_nb) {
Ok(()) => Ok(StripSuccess::Stripped),
#[cfg(not(tarpaulin_include))]
Err(e) => Err(e.into()),
}
}
}
}
pub fn write_nb<W, T>(mut writer: W, value: &T) -> Result<(), NBWriteError>
where
W: Write,
T: ?Sized + Serialize,
{
let formatter = serde_json::ser::PrettyFormatter::with_indent(b" ");
let mut ser = serde_json::Serializer::with_formatter(&mut writer, formatter);
value.serialize(&mut ser)?;
writeln!(writer)?;
Ok(())
}
#[derive(Debug, Error)]
pub enum StripError {
#[error("File read Error")]
ReadError(#[from] NBReadError),
#[error("File write Error")]
WriteError(#[from] NBWriteError),
}
#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
pub enum StripSuccess {
NoChange,
Stripped,
}
#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
pub enum StripResult {
NoChange,
Stripped,
ReadError(String),
WriteError(String),
}
impl From<StripSuccess> for StripResult {
fn from(value: StripSuccess) -> Self {
match value {
StripSuccess::NoChange => Self::NoChange,
StripSuccess::Stripped => Self::Stripped,
}
}
}
impl From<StripError> for StripResult {
fn from(value: StripError) -> Self {
match value {
StripError::ReadError(e) => Self::ReadError(e.to_string()),
StripError::WriteError(e) => Self::WriteError(e.to_string()),
}
}
}
impl From<Result<StripSuccess, StripError>> for StripResult {
fn from(value: Result<StripSuccess, StripError>) -> Self {
match value {
Ok(v) => v.into(),
Err(v) => v.into(),
}
}
}
impl Display for StripResult {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NoChange => write!(f, "No Change"),
Self::Stripped => write!(f, "Stripped"),
Self::ReadError(e) => write!(f, "Read error: {e}"),
Self::WriteError(e) => write!(f, "Write error: {e}"),
}
}
}
impl StripSuccess {
pub const fn from_stripped(stripped: bool) -> Self {
if stripped {
Self::Stripped
} else {
Self::NoChange
}
}
}
impl StripResult {
pub const fn is_err(&self) -> bool {
matches!(self, Self::ReadError(_) | Self::WriteError(_))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_strip_error_to_strip_result() {
let read_error = StripError::ReadError(NBReadError::IO(std::io::Error::from(
std::io::ErrorKind::NotFound,
)));
let write_error = StripError::WriteError(NBWriteError::IO(std::io::Error::from(
std::io::ErrorKind::PermissionDenied,
)));
let read_error_res: StripResult = read_error.into();
let write_error_res: StripResult = write_error.into();
assert!(matches!(read_error_res, StripResult::ReadError(..)));
assert!(matches!(write_error_res, StripResult::WriteError(..)));
assert!(write_error_res.to_string().starts_with("Write error:"));
}
}