Skip to main content

entrenar/research/notebook/
cell.rs

1//! Cell types for Jupyter notebooks.
2
3use serde::{Deserialize, Serialize};
4
5use super::split_source;
6
7/// Cell type in a notebook
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
9#[serde(rename_all = "snake_case")]
10pub enum CellType {
11    /// Executable code cell
12    Code,
13    /// Markdown documentation cell
14    Markdown,
15    /// Raw text cell
16    Raw,
17}
18
19impl std::fmt::Display for CellType {
20    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        match self {
22            Self::Code => write!(f, "code"),
23            Self::Markdown => write!(f, "markdown"),
24            Self::Raw => write!(f, "raw"),
25        }
26    }
27}
28
29/// A single cell in a notebook
30#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
31pub struct Cell {
32    /// Type of cell
33    pub cell_type: CellType,
34    /// Source content (lines)
35    pub source: Vec<String>,
36    /// Cell outputs (for code cells)
37    pub outputs: Vec<CellOutput>,
38    /// Execution count (for code cells)
39    pub execution_count: Option<u32>,
40    /// Cell metadata
41    pub metadata: CellMetadata,
42}
43
44impl Cell {
45    /// Create a new code cell
46    pub fn code(source: impl Into<String>) -> Self {
47        Self {
48            cell_type: CellType::Code,
49            source: split_source(source.into()),
50            outputs: Vec::new(),
51            execution_count: None,
52            metadata: CellMetadata::default(),
53        }
54    }
55
56    /// Create a new markdown cell
57    pub fn markdown(source: impl Into<String>) -> Self {
58        Self {
59            cell_type: CellType::Markdown,
60            source: split_source(source.into()),
61            outputs: Vec::new(),
62            execution_count: None,
63            metadata: CellMetadata::default(),
64        }
65    }
66
67    /// Create a new raw cell
68    pub fn raw(source: impl Into<String>) -> Self {
69        Self {
70            cell_type: CellType::Raw,
71            source: split_source(source.into()),
72            outputs: Vec::new(),
73            execution_count: None,
74            metadata: CellMetadata::default(),
75        }
76    }
77
78    /// Add an output to the cell
79    pub fn with_output(mut self, output: CellOutput) -> Self {
80        self.outputs.push(output);
81        self
82    }
83
84    /// Set execution count
85    pub fn with_execution_count(mut self, count: u32) -> Self {
86        self.execution_count = Some(count);
87        self
88    }
89
90    /// Get the source as a single string
91    pub fn source_text(&self) -> String {
92        self.source.join("")
93    }
94}
95
96/// Cell output
97#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
98pub struct CellOutput {
99    /// Output type
100    pub output_type: String,
101    /// Output data
102    pub data: Option<serde_json::Value>,
103    /// Text output
104    pub text: Option<Vec<String>>,
105}
106
107impl CellOutput {
108    /// Create a stream output (stdout/stderr)
109    pub fn stream(_name: &str, text: impl Into<String>) -> Self {
110        Self {
111            output_type: "stream".to_string(),
112            data: None,
113            text: Some(split_source(text.into())),
114        }
115    }
116
117    /// Create an execute result output
118    pub fn execute_result(data: serde_json::Value) -> Self {
119        Self { output_type: "execute_result".to_string(), data: Some(data), text: None }
120    }
121}
122
123/// Cell metadata
124#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
125pub struct CellMetadata {
126    /// Whether the cell is collapsed
127    #[serde(skip_serializing_if = "Option::is_none")]
128    pub collapsed: Option<bool>,
129    /// Whether the cell is scrolled
130    #[serde(skip_serializing_if = "Option::is_none")]
131    pub scrolled: Option<bool>,
132}