use std::path::Path;
use serde_json::Value;
pub fn is_notebook(path: &Path) -> bool {
path.extension()
.and_then(|ext| ext.to_str())
.is_some_and(|ext| ext.eq_ignore_ascii_case("ipynb"))
}
pub fn extract_notebook_content(json_text: &str) -> Option<String> {
let json: Value = serde_json::from_str(json_text).ok()?;
let cells = json.get("cells")?.as_array()?;
let mut parts: Vec<String> = Vec::new();
for cell in cells {
if let Some(source) = cell.get("source").and_then(collect_text) {
parts.push(source);
}
if let Some(outputs) = cell.get("outputs").and_then(Value::as_array) {
for output in outputs {
if let Some(text) = output.get("text").and_then(collect_text) {
parts.push(text);
}
if let Some(text) = output
.get("data")
.and_then(|data| data.get("text/plain"))
.and_then(collect_text)
{
parts.push(text);
}
}
}
}
if parts.is_empty() {
None
} else {
Some(
parts
.iter()
.map(|part| part.trim_end_matches('\n'))
.collect::<Vec<_>>()
.join("\n"),
)
}
}
fn collect_text(value: &Value) -> Option<String> {
match value {
Value::String(s) => (!s.is_empty()).then(|| s.clone()),
Value::Array(items) => {
let combined: String = items.iter().filter_map(Value::as_str).collect();
(!combined.is_empty()).then_some(combined)
}
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
#[test]
fn test_is_notebook() {
assert!(is_notebook(&PathBuf::from("Analysis.ipynb")));
assert!(is_notebook(&PathBuf::from("a/b/NOTEBOOK.IPYNB")));
assert!(!is_notebook(&PathBuf::from("script.py")));
assert!(!is_notebook(&PathBuf::from("data.json")));
}
#[test]
fn test_extract_source_array_and_stream_output() {
let json = r#"{
"cells": [
{"cell_type":"code","source":["@show typeof(C)\n","C[1:10,:]\n"],
"outputs":[{"output_type":"stream","name":"stdout",
"text":["\t(c) Brendan O'Donoghue, Stanford University, 2012\n"]}]}
]
}"#;
let text = extract_notebook_content(json).expect("should extract");
assert!(text.contains("@show typeof(C)\nC[1:10,:]"));
assert!(text.contains("(c) Brendan O'Donoghue, Stanford University, 2012"));
assert!(!text.contains("\", \""));
}
#[test]
fn test_extract_skips_binary_output_data() {
let json = r#"{
"cells": [
{"cell_type":"code","source":"print(1)",
"outputs":[{"output_type":"display_data",
"data":{"image/png":"iVBORw0KGgoAAAANSU","text/plain":"<Figure>"}}]}
]
}"#;
let text = extract_notebook_content(json).expect("should extract");
assert!(text.contains("print(1)"));
assert!(text.contains("<Figure>"));
assert!(!text.contains("iVBORw0KGgoAAAANSU"));
}
#[test]
fn test_extract_source_string_form() {
let json = r##"{"cells":[{"cell_type":"markdown","source":"# Title\nSome prose"}]}"##;
let text = extract_notebook_content(json).expect("should extract");
assert_eq!(text, "# Title\nSome prose");
}
#[test]
fn test_extract_invalid_or_empty_returns_none() {
assert!(extract_notebook_content("not json").is_none());
assert!(extract_notebook_content(r#"{"nbformat":4}"#).is_none());
assert!(extract_notebook_content(r#"{"cells":[]}"#).is_none());
}
#[test]
fn test_extract_no_blank_line_between_parts() {
let json = r#"{
"cells": [
{"cell_type":"code","source":"a()","outputs":[{"output_type":"stream","text":["one\n","two\n"]}]},
{"cell_type":"code","source":"b()"}
]
}"#;
let text = extract_notebook_content(json).expect("should extract");
assert!(
!text.contains("\n\n"),
"no blank lines between parts: {text:?}"
);
assert!(text.contains("two\nb()"));
}
}