use serde::{Deserialize, Serialize};
use serde_json::json;
use super::cell::{Cell, CellType};
use super::kernel::KernelSpec;
use crate::research::literate::LiterateDocument;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NotebookExporter {
pub cells: Vec<Cell>,
pub kernel: KernelSpec,
pub metadata: NotebookMetadata,
}
impl Default for NotebookExporter {
fn default() -> Self {
Self::new()
}
}
impl NotebookExporter {
pub fn new() -> Self {
Self {
cells: Vec::new(),
kernel: KernelSpec::python3(),
metadata: NotebookMetadata::default(),
}
}
pub fn with_kernel(kernel: KernelSpec) -> Self {
Self { cells: Vec::new(), kernel, metadata: NotebookMetadata::default() }
}
pub fn add_cell(&mut self, cell: Cell) {
self.cells.push(cell);
}
pub fn add_code(&mut self, source: impl Into<String>) {
self.cells.push(Cell::code(source));
}
pub fn add_markdown(&mut self, source: impl Into<String>) {
self.cells.push(Cell::markdown(source));
}
pub fn from_literate(doc: &LiterateDocument) -> Self {
let mut exporter = Self::new();
if doc.is_typst() || doc.is_markdown() {
let content = doc.content();
let blocks = doc.extract_code_blocks();
let primary_lang = blocks.iter().find_map(|b| b.language.as_ref()).map(String::as_str);
exporter.kernel = match primary_lang {
Some("rust") => KernelSpec::evcxr(),
Some("julia") => KernelSpec::julia(),
Some(other_lang) => {
eprintln!(
"Warning: unsupported kernel language '{other_lang}', defaulting to Python 3"
);
KernelSpec::python3()
}
None => KernelSpec::python3(),
};
let mut last_end = 0;
for block in &blocks {
let block_pattern = format!("```{}", block.language.as_deref().unwrap_or(""));
if let Some(start_pos) = content[last_end..].find(&block_pattern) {
let absolute_start = last_end + start_pos;
let markdown_content = &content[last_end..absolute_start];
let trimmed = markdown_content.trim();
if !trimmed.is_empty() {
exporter.add_markdown(trimmed);
}
exporter.add_code(&block.content);
let code_end = content[absolute_start..]
.find("```\n")
.or_else(|| content[absolute_start..].find("```"))
.map_or(content.len(), |p| {
absolute_start
+ p
+ content[absolute_start + p..].find('\n').unwrap_or(3)
+ 1
});
last_end = code_end.min(content.len());
}
}
if last_end < content.len() {
let remaining = content[last_end..].trim();
if !remaining.is_empty() {
exporter.add_markdown(remaining);
}
}
} else {
exporter.add_markdown(doc.content());
}
exporter
}
pub fn to_ipynb(&self) -> String {
let notebook = json!({
"nbformat": 4,
"nbformat_minor": 5,
"metadata": {
"kernelspec": {
"display_name": self.kernel.display_name,
"language": self.kernel.language,
"name": self.kernel.name
},
"language_info": {
"name": self.kernel.language
}
},
"cells": self.cells.iter().map(|cell| {
let mut cell_json = json!({
"cell_type": cell.cell_type.to_string(),
"source": cell.source,
"metadata": cell.metadata
});
if cell.cell_type == CellType::Code {
cell_json["outputs"] = json!(cell.outputs.iter().map(|o| {
let mut out = json!({
"output_type": o.output_type
});
if let Some(data) = &o.data {
out["data"] = data.clone();
}
if let Some(text) = &o.text {
out["text"] = json!(text);
out["name"] = json!("stdout");
}
out
}).collect::<Vec<_>>());
cell_json["execution_count"] = json!(cell.execution_count);
}
cell_json
}).collect::<Vec<_>>()
});
serde_json::to_string_pretty(¬ebook).unwrap_or_else(|_err| "{}".to_string())
}
pub fn cell_count(&self) -> usize {
self.cells.len()
}
pub fn code_cells(&self) -> Vec<&Cell> {
self.cells.iter().filter(|c| c.cell_type == CellType::Code).collect()
}
pub fn markdown_cells(&self) -> Vec<&Cell> {
self.cells.iter().filter(|c| c.cell_type == CellType::Markdown).collect()
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct NotebookMetadata {
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
#[serde(skip_serializing_if = "Vec::is_empty", default)]
pub authors: Vec<String>,
}
#[cfg(test)]
mod tests {
use super::super::cell::CellOutput;
use super::*;
#[test]
fn test_kernel_selection_all_language_variants() {
let languages: &[Option<&str>] = &[Some("rust"), Some("julia"), Some("javascript"), None];
for lang in languages {
let kernel = match *lang {
Some("rust") => KernelSpec::evcxr(),
Some("julia") => KernelSpec::julia(),
Some(_other_lang) => KernelSpec::python3(),
None => KernelSpec::python3(),
};
match lang {
Some("rust") => assert_eq!(kernel.language, "rust"),
Some("julia") => assert_eq!(kernel.language, "julia"),
Some(_) => assert_eq!(kernel.language, "python"),
None => assert_eq!(kernel.language, "python"),
}
}
}
#[test]
fn test_notebook_exporter_new() {
let exporter = NotebookExporter::new();
assert_eq!(exporter.cells.len(), 0);
assert_eq!(exporter.kernel.language, "python");
}
#[test]
fn test_notebook_exporter_default() {
let exporter = NotebookExporter::default();
assert_eq!(exporter.cell_count(), 0);
}
#[test]
fn test_add_code_and_markdown() {
let mut exporter = NotebookExporter::new();
exporter.add_code("print('hello')");
exporter.add_markdown("# Title");
assert_eq!(exporter.cell_count(), 2);
assert_eq!(exporter.code_cells().len(), 1);
assert_eq!(exporter.markdown_cells().len(), 1);
}
#[test]
fn test_notebook_metadata_default() {
let meta = NotebookMetadata::default();
assert!(meta.title.is_none());
assert!(meta.authors.is_empty());
}
#[test]
fn test_notebook_exporter_with_kernel() {
let exporter = NotebookExporter::with_kernel(KernelSpec::evcxr());
assert_eq!(exporter.kernel.language, "rust");
assert_eq!(exporter.cell_count(), 0);
}
#[test]
fn test_notebook_exporter_with_julia_kernel() {
let exporter = NotebookExporter::with_kernel(KernelSpec::julia());
assert_eq!(exporter.kernel.language, "julia");
assert_eq!(exporter.kernel.display_name, "Julia 1.9");
}
#[test]
fn test_add_cell_directly() {
let mut exporter = NotebookExporter::new();
exporter.add_cell(Cell::code("x = 1"));
exporter.add_cell(Cell::markdown("# Hello"));
exporter.add_cell(Cell::raw("raw text"));
assert_eq!(exporter.cell_count(), 3);
assert_eq!(exporter.code_cells().len(), 1);
assert_eq!(exporter.markdown_cells().len(), 1);
}
#[test]
fn test_to_ipynb_empty() {
let exporter = NotebookExporter::new();
let json = exporter.to_ipynb();
assert!(json.contains("nbformat"));
assert!(json.contains("\"cells\": []"));
}
#[test]
fn test_to_ipynb_with_cells() {
let mut exporter = NotebookExporter::new();
exporter.add_code("print('hello')");
exporter.add_markdown("# Title");
let json = exporter.to_ipynb();
assert!(json.contains("print('hello')"));
assert!(json.contains("# Title"));
assert!(json.contains("\"cell_type\": \"code\""));
assert!(json.contains("\"cell_type\": \"markdown\""));
}
#[test]
fn test_to_ipynb_kernelspec() {
let exporter = NotebookExporter::with_kernel(KernelSpec::evcxr());
let json = exporter.to_ipynb();
assert!(json.contains("\"language\": \"rust\""));
assert!(json.contains("Rust"));
}
#[test]
fn test_to_ipynb_code_cell_has_outputs_and_execution_count() {
let mut exporter = NotebookExporter::new();
exporter.add_code("1 + 1");
let json = exporter.to_ipynb();
assert!(json.contains("\"outputs\""));
assert!(json.contains("\"execution_count\""));
}
#[test]
fn test_to_ipynb_with_output() {
let mut exporter = NotebookExporter::new();
let cell = Cell::code("print(42)")
.with_output(CellOutput::stream("stdout", "42\n"))
.with_execution_count(1);
exporter.add_cell(cell);
let json = exporter.to_ipynb();
assert!(json.contains("stream"));
assert!(json.contains("42"));
assert!(json.contains("stdout"));
}
#[test]
fn test_to_ipynb_with_execute_result() {
let mut exporter = NotebookExporter::new();
let data = serde_json::json!({"text/plain": ["result"]});
let cell = Cell::code("1 + 1").with_output(CellOutput::execute_result(data));
exporter.add_cell(cell);
let json = exporter.to_ipynb();
assert!(json.contains("execute_result"));
assert!(json.contains("text/plain"));
}
#[test]
fn test_from_literate_markdown() {
let doc = LiterateDocument::parse_markdown(
"# Hello\n\nSome text.\n\n```python\nprint('hi')\n```\n\nMore text.",
);
let exporter = NotebookExporter::from_literate(&doc);
assert!(exporter.cell_count() > 0);
assert!(!exporter.code_cells().is_empty());
assert!(!exporter.markdown_cells().is_empty());
assert_eq!(exporter.kernel.language, "python");
}
#[test]
fn test_from_literate_rust_kernel() {
let doc =
LiterateDocument::parse_markdown("# Rust Example\n\n```rust\nfn main() {}\n```\n");
let exporter = NotebookExporter::from_literate(&doc);
assert_eq!(exporter.kernel.language, "rust");
}
#[test]
fn test_from_literate_julia_kernel() {
let doc =
LiterateDocument::parse_markdown("# Julia Example\n\n```julia\nprintln(\"hi\")\n```\n");
let exporter = NotebookExporter::from_literate(&doc);
assert_eq!(exporter.kernel.language, "julia");
}
#[test]
fn test_from_literate_no_code_blocks() {
let doc = LiterateDocument::parse_markdown("# Just Markdown\n\nNo code here.");
let exporter = NotebookExporter::from_literate(&doc);
assert!(exporter.cell_count() >= 1);
assert!(exporter.code_cells().is_empty());
}
#[test]
fn test_from_literate_raw_text() {
let doc = LiterateDocument::raw("Just plain text, no special parsing.");
let exporter = NotebookExporter::from_literate(&doc);
assert_eq!(exporter.cell_count(), 1);
assert_eq!(exporter.markdown_cells().len(), 1);
}
#[test]
fn test_from_literate_multiple_code_blocks() {
let doc = LiterateDocument::parse_markdown(
"Intro\n\n```python\nx = 1\n```\n\nMiddle text\n\n```python\ny = 2\n```\n\nEnd",
);
let exporter = NotebookExporter::from_literate(&doc);
assert_eq!(exporter.code_cells().len(), 2);
}
#[test]
fn test_notebook_metadata_with_values() {
let meta = NotebookMetadata {
title: Some("My Notebook".to_string()),
authors: vec!["Author A".to_string(), "Author B".to_string()],
};
assert_eq!(meta.title.as_deref(), Some("My Notebook"));
assert_eq!(meta.authors.len(), 2);
}
#[test]
fn test_notebook_metadata_serialization() {
let meta = NotebookMetadata {
title: Some("Test".to_string()),
authors: vec!["Alice".to_string()],
};
let json = serde_json::to_string(&meta).expect("serialize");
let restored: NotebookMetadata = serde_json::from_str(&json).expect("deserialize");
assert_eq!(restored.title, meta.title);
assert_eq!(restored.authors, meta.authors);
}
#[test]
fn test_notebook_metadata_serialization_skip_empty() {
let meta = NotebookMetadata::default();
let json = serde_json::to_string(&meta).expect("serialize");
assert!(!json.contains("title"));
assert!(!json.contains("authors"));
}
#[test]
fn test_notebook_exporter_serialization() {
let mut exporter = NotebookExporter::new();
exporter.add_code("x = 1");
let json = serde_json::to_string(&exporter).expect("serialize");
let restored: NotebookExporter = serde_json::from_str(&json).expect("deserialize");
assert_eq!(restored.cell_count(), 1);
assert_eq!(restored.kernel.language, "python");
}
#[test]
fn test_code_cells_and_markdown_cells_filtering() {
let mut exporter = NotebookExporter::new();
exporter.add_code("a");
exporter.add_markdown("b");
exporter.add_code("c");
exporter.add_markdown("d");
exporter.add_code("e");
assert_eq!(exporter.code_cells().len(), 3);
assert_eq!(exporter.markdown_cells().len(), 2);
assert_eq!(exporter.cell_count(), 5);
}
#[test]
fn test_from_literate_typst() {
let doc = LiterateDocument::Typst(
"= Title\n\nSome text.\n\n```python\nprint('hi')\n```\n\nMore text.".to_string(),
);
let exporter = NotebookExporter::from_literate(&doc);
assert!(exporter.cell_count() > 0);
}
}