jtool_notebook/

notebook.rs

use crate::cell::Cell;
use crate::error::NotebookError;
use crate::timing::{CellTiming, TimingSummary};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use std::str::FromStr;

/// A Jupyter notebook
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Notebook {
    /// Notebook cells
    pub cells: Vec<Cell>,
    /// Notebook-level metadata
    pub metadata: NotebookMetadata,
    /// nbformat version (major)
    pub nbformat: u32,
    /// nbformat version (minor)
    pub nbformat_minor: u32,
}

impl Notebook {
    /// Read a notebook from a file
    pub fn from_file(path: &Path) -> Result<Self, NotebookError> {
        let contents = fs::read_to_string(path)?;
        contents.parse()
    }

    /// Write a notebook to a file
    pub fn to_file(&self, path: &Path) -> Result<(), NotebookError> {
        let contents = serde_json::to_string_pretty(self)?;
        fs::write(path, contents)?;
        Ok(())
    }

    /// Convert notebook to JSON string
    pub fn to_string(&self) -> Result<String, NotebookError> {
        Ok(serde_json::to_string_pretty(self)?)
    }

    /// Get a cell by index
    pub fn get_cell(&self, index: usize) -> Result<&Cell, NotebookError> {
        self.cells
            .get(index)
            .ok_or(NotebookError::CellIndexOutOfBounds(index, self.cells.len()))
    }

    /// Get a mutable cell by index
    pub fn get_cell_mut(&mut self, index: usize) -> Result<&mut Cell, NotebookError> {
        let len = self.cells.len();
        self.cells
            .get_mut(index)
            .ok_or(NotebookError::CellIndexOutOfBounds(index, len))
    }

    /// Get multiple cells by indices
    pub fn get_cells(&self, indices: &[usize]) -> Result<Vec<&Cell>, NotebookError> {
        indices.iter().map(|&i| self.get_cell(i)).collect()
    }

    /// Count cells
    pub fn cell_count(&self) -> usize {
        self.cells.len()
    }

    /// Count code cells
    pub fn code_cell_count(&self) -> usize {
        self.cells.iter().filter(|c| c.is_code()).count()
    }

    /// Get timing information for all cells
    ///
    /// Returns a vector of CellTiming for all code cells that have timing metadata.
    /// Cells without timing metadata are skipped.
    pub fn cell_timings(&self) -> Vec<CellTiming> {
        self.cells
            .iter()
            .enumerate()
            .filter_map(|(i, cell)| {
                if let Cell::Code(code_cell) = cell {
                    let timing = code_cell.timing()?;
                    Some(CellTiming::new(i, code_cell.execution_count, timing))
                } else {
                    None
                }
            })
            .collect()
    }

    /// Get timing summary for the notebook
    ///
    /// Returns timing statistics (total, average, min, max) for all cells with timing data.
    /// Returns None if no cells have timing metadata.
    pub fn timing_summary(&self) -> Option<TimingSummary> {
        let timings = self.cell_timings();
        TimingSummary::from_cells(&timings)
    }

    /// Count cells with timing metadata
    pub fn cells_with_timing_count(&self) -> usize {
        self.cells
            .iter()
            .filter(|cell| {
                if let Cell::Code(code_cell) = cell {
                    code_cell.has_timing()
                } else {
                    false
                }
            })
            .count()
    }
}

impl FromStr for Notebook {
    type Err = NotebookError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let notebook: Notebook = serde_json::from_str(s)?;

        // Validate nbformat version
        if notebook.nbformat < 4 {
            return Err(NotebookError::UnsupportedVersion(
                notebook.nbformat,
                notebook.nbformat_minor,
            ));
        }

        Ok(notebook)
    }
}

/// Notebook metadata
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct NotebookMetadata {
    /// Kernel specification
    #[serde(skip_serializing_if = "Option::is_none")]
    pub kernelspec: Option<KernelSpec>,
    /// Language info
    #[serde(skip_serializing_if = "Option::is_none")]
    pub language_info: Option<LanguageInfo>,
    /// Additional metadata (extensible)
    #[serde(flatten)]
    pub extra: HashMap<String, serde_json::Value>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KernelSpec {
    pub name: String,
    pub display_name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub language: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LanguageInfo {
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub version: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub file_extension: Option<String>,
    #[serde(flatten)]
    pub extra: HashMap<String, serde_json::Value>,
}

#[cfg(test)]
mod tests {
    use super::*;

    fn create_minimal_notebook() -> Notebook {
        Notebook {
            cells: vec![],
            metadata: NotebookMetadata::default(),
            nbformat: 4,
            nbformat_minor: 5,
        }
    }

    fn create_test_notebook_json() -> &'static str {
        r#"{
            "cells": [
                {
                    "cell_type": "code",
                    "execution_count": null,
                    "id": "test-1",
                    "metadata": {},
                    "outputs": [],
                    "source": ["print('hello')"]
                }
            ],
            "metadata": {},
            "nbformat": 4,
            "nbformat_minor": 5
        }"#
    }

    #[test]
    fn test_parse_valid_notebook() -> Result<(), Box<dyn std::error::Error>> {
        let json = create_test_notebook_json();
        let notebook: Notebook = json.parse()?;

        assert_eq!(notebook.nbformat, 4);
        assert_eq!(notebook.nbformat_minor, 5);
        assert_eq!(notebook.cells.len(), 1);
        Ok(())
    }

    #[test]
    fn test_parse_notebook_with_missing_optional_fields() {
        let json = r#"{
            "cells": [{
                "cell_type": "code",
                "id": "test",
                "metadata": {},
                "source": []
            }],
            "metadata": {},
            "nbformat": 4,
            "nbformat_minor": 5
        }"#;

        let result: Result<Notebook, _> = json.parse();
        // Should handle missing optional fields like execution_count, outputs
        assert!(result.is_ok());
    }

    #[test]
    fn test_parse_malformed_json() {
        let json = "{ invalid json }";
        let result: Result<Notebook, _> = json.parse();
        assert!(result.is_err());
    }

    #[test]
    fn test_unsupported_version() {
        let json = r#"{
            "cells": [],
            "metadata": {},
            "nbformat": 3,
            "nbformat_minor": 0
        }"#;

        let result: Result<Notebook, _> = json.parse();
        assert!(result.is_err());

        assert!(matches!(
            result,
            Err(NotebookError::UnsupportedVersion(3, 0))
        ));
    }

    #[test]
    fn test_serialize_notebook() -> Result<(), Box<dyn std::error::Error>> {
        let notebook = create_minimal_notebook();
        let json = notebook.to_string()?;

        // Should be valid JSON
        let parsed: serde_json::Value = serde_json::from_str(&json)?;
        assert_eq!(parsed["nbformat"], 4);
        assert_eq!(parsed["nbformat_minor"], 5);
        Ok(())
    }

    #[test]
    fn test_roundtrip_serialization() -> Result<(), Box<dyn std::error::Error>> {
        let json = create_test_notebook_json();
        let notebook1: Notebook = json.parse()?;
        let serialized = notebook1.to_string()?;
        let notebook2: Notebook = serialized.parse()?;

        assert_eq!(notebook1.cells.len(), notebook2.cells.len());
        assert_eq!(notebook1.nbformat, notebook2.nbformat);
        Ok(())
    }

    #[test]
    fn test_get_cell() -> Result<(), Box<dyn std::error::Error>> {
        let json = create_test_notebook_json();
        let notebook: Notebook = json.parse()?;

        let cell = notebook.get_cell(0)?;
        assert!(matches!(cell, Cell::Code(_)));
        Ok(())
    }

    #[test]
    fn test_get_cell_out_of_bounds() {
        let notebook = create_minimal_notebook();
        let result = notebook.get_cell(0);

        assert!(result.is_err());
        assert!(matches!(
            result,
            Err(NotebookError::CellIndexOutOfBounds(0, 0))
        ));
    }

    #[test]
    fn test_cell_count() -> Result<(), Box<dyn std::error::Error>> {
        let notebook = create_minimal_notebook();
        assert_eq!(notebook.cell_count(), 0);

        let json = create_test_notebook_json();
        let notebook: Notebook = json.parse()?;
        assert_eq!(notebook.cell_count(), 1);
        Ok(())
    }

    #[test]
    fn test_code_cell_count() -> Result<(), Box<dyn std::error::Error>> {
        let json = r#"{
            "cells": [
                {
                    "cell_type": "code",
                    "id": "1",
                    "metadata": {},
                    "source": []
                },
                {
                    "cell_type": "markdown",
                    "id": "2",
                    "metadata": {},
                    "source": []
                },
                {
                    "cell_type": "code",
                    "id": "3",
                    "metadata": {},
                    "source": []
                }
            ],
            "metadata": {},
            "nbformat": 4,
            "nbformat_minor": 5
        }"#;

        let notebook: Notebook = json.parse()?;
        assert_eq!(notebook.code_cell_count(), 2);
        assert_eq!(notebook.cell_count(), 3);
        Ok(())
    }

    #[test]
    fn test_preserve_cell_metadata() -> Result<(), Box<dyn std::error::Error>> {
        let json = r#"{
            "cells": [{
                "cell_type": "code",
                "id": "test",
                "metadata": {
                    "tags": ["important"],
                    "custom": "value"
                },
                "source": []
            }],
            "metadata": {},
            "nbformat": 4,
            "nbformat_minor": 5
        }"#;

        let notebook: Notebook = json.parse()?;
        let serialized = notebook.to_string()?;

        // Metadata should be preserved
        assert!(serialized.contains("important"));
        assert!(serialized.contains("custom"));
        Ok(())
    }

    #[test]
    fn test_empty_notebook() {
        let notebook = create_minimal_notebook();
        assert_eq!(notebook.cells.len(), 0);
        assert_eq!(notebook.code_cell_count(), 0);
    }

    #[test]
    fn test_notebook_with_kernel_metadata() -> Result<(), Box<dyn std::error::Error>> {
        let json = r#"{
            "cells": [],
            "metadata": {
                "kernelspec": {
                    "name": "python3",
                    "display_name": "Python 3",
                    "language": "python"
                },
                "language_info": {
                    "name": "python",
                    "version": "3.11.0"
                }
            },
            "nbformat": 4,
            "nbformat_minor": 5
        }"#;

        let notebook: Notebook = json.parse()?;
        assert!(notebook.metadata.kernelspec.is_some());
        assert!(notebook.metadata.language_info.is_some());

        if let Some(kernelspec) = notebook.metadata.kernelspec {
            assert_eq!(kernelspec.name, "python3");
            assert_eq!(kernelspec.display_name, "Python 3");
        }
        Ok(())
    }

    #[test]
    fn test_negative_execution_count() -> Result<(), Box<dyn std::error::Error>> {
        let json = r#"{
            "cells": [
                {
                    "cell_type": "code",
                    "execution_count": -1,
                    "id": "test-negative",
                    "metadata": {},
                    "outputs": [],
                    "source": ["print('test')"]
                }
            ],
            "metadata": {},
            "nbformat": 4,
            "nbformat_minor": 5
        }"#;

        let notebook: Notebook = json.parse()?;
        assert_eq!(notebook.cells.len(), 1);

        assert!(
            matches!(
                &notebook.cells[0],
                Cell::Code(code_cell) if code_cell.execution_count == Some(-1)
            ),
            "Expected code cell with execution_count -1"
        );
        Ok(())
    }
}