use crate::extra_keys::ExtraKey;
use crate::schema::{Cell, RawNotebook};
use itertools::Itertools;
use serde_json::Value;
pub fn pop_value_child<T: AsRef<str>>(value: &mut serde_json::Value, path: &[T]) -> Option<Value> {
let n_parts = path.len();
let trial_key = path.iter().map(std::convert::AsRef::as_ref).join(".");
let removed = value
.as_object_mut()
.and_then(|x| x.remove(trial_key.as_str()));
if removed.is_some() {
return removed;
}
for i in (1..n_parts).rev() {
let trial_key = path[0..i].iter().map(std::convert::AsRef::as_ref).join(".");
let tail_key = &path[i..n_parts];
#[allow(clippy::unwrap_used)]
if value
.as_object()
.is_some_and(|x| x.contains_key(trial_key.as_str()))
{
let inner = value
.as_object_mut()
.and_then(|x| x.get_mut(trial_key.as_str()))
.unwrap();
return pop_value_child(inner, tail_key);
}
}
None
}
pub fn pop_cell_key(cell: &mut Cell, extra_key: &ExtraKey) -> Option<serde_json::Value> {
let (cell, ExtraKey::CellMeta(cellmeta_key)) = (cell, extra_key) else {
return None;
};
pop_value_child(cell.get_metadata_mut(), cellmeta_key.parts.as_slice())
}
pub fn pop_meta_key(nb: &mut RawNotebook, extra_key: &ExtraKey) -> Option<serde_json::Value> {
let (nb, ExtraKey::Metadata(meta_key)) = (nb, extra_key) else {
return None;
};
pop_value_child(&mut nb.metadata, meta_key.parts.as_slice())
}
pub fn get_value_child<'a, T: AsRef<str>>(value: &'a Value, path: &[T]) -> Option<&'a Value> {
let n_parts = path.len();
if n_parts == 0 {
return Some(value);
}
for i in (1..=n_parts).rev() {
let trial_key = path[0..i].iter().map(std::convert::AsRef::as_ref).join(".");
let tail_key = &path[i..n_parts];
if let Some(inner) = value.as_object().and_then(|x| x.get(trial_key.as_str())) {
return get_value_child(inner, tail_key);
}
}
None
}
#[allow(clippy::unwrap_used)]
#[cfg(test)]
mod test {
use std::str::FromStr;
use crate::schema::CodeCell;
use crate::schema::SourceValue;
use super::*;
use serde_json::json;
#[test]
fn test_get_value_child_simple() {
let metadata = json!({"keep_output": true});
let keep_output = get_value_child(&metadata, &["keep_output"]);
assert!(keep_output.is_some_and(|x| x.as_bool().is_some_and(|b| b)));
}
#[test]
fn test_get_value_child_period() {
let mut metadata = json!({"application/vnd.databricks.v1+cell": "bananas"});
let key = ExtraKey::from_str("cell.metadata.application/vnd.databricks.v1+cell").unwrap();
println!("{:?}", key.get_parts());
println!("{key:?}");
let output = get_value_child(&metadata, key.get_parts());
println!("{output:?}");
assert!(output.is_some());
assert_eq!(output.unwrap().as_str().unwrap(), "bananas");
pop_value_child(&mut metadata, key.get_parts());
assert!(metadata.as_object().is_some_and(serde_json::Map::is_empty));
}
#[test]
fn test_pop_nothing() {
let mut cell = Cell::Code(CodeCell {
execution_count: None,
id: None,
metadata: json!({"hello": "world"}),
outputs: vec![],
source: SourceValue::StringArray(vec![]),
});
let nothing = pop_cell_key(&mut cell, &ExtraKey::from_str("metadata.hello").unwrap());
assert!(nothing.is_none());
let mut nb = RawNotebook {
cells: vec![],
nbformat: 4,
nbformat_minor: 5,
metadata: json!({"hello": "world"}),
};
let nothing = pop_meta_key(&mut nb, &ExtraKey::from_str("cell.metadata.hello").unwrap());
assert!(nothing.is_none());
}
}