kernel_sidecar/handlers/
outputs.rs

1use tokio::sync::Mutex;
2
3use crate::handlers::Handler;
4use crate::jupyter::response::Response;
5use crate::notebook::{Notebook, Output};
6
7use std::fmt::Debug;
8use std::sync::Arc;
9
10// Update a document model with outputs while running a cell
11#[derive(Debug)]
12pub struct OutputHandler {
13    nb: Arc<Mutex<Notebook>>,
14    cell_id: String,
15    clear_on_next_output: bool,
16}
17
18impl OutputHandler {
19    pub fn new(nb: Arc<Mutex<Notebook>>, cell_id: &str) -> Self {
20        Self {
21            nb,
22            cell_id: cell_id.to_string(),
23            clear_on_next_output: false,
24        }
25    }
26
27    pub async fn add_output(&mut self, content: Output) {
28        let mut nb = self.nb.lock().await;
29        if let Some(cell) = nb.get_mut_cell(&self.cell_id) {
30            cell.add_output(content);
31        }
32    }
33
34    pub async fn clear_output(&mut self) {
35        let mut nb = self.nb.lock().await;
36        if let Some(cell) = nb.get_mut_cell(&self.cell_id) {
37            cell.clear_output();
38        }
39    }
40}
41
42#[async_trait::async_trait]
43impl Handler for OutputHandler {
44    async fn handle(&mut self, msg: &Response) {
45        match msg {
46            Response::ExecuteResult(m) => {
47                let output = Output::ExecuteResult(m.content.clone());
48                if self.clear_on_next_output {
49                    self.clear_output().await;
50                    self.clear_on_next_output = false;
51                }
52                self.add_output(output).await;
53            }
54            Response::Stream(m) => {
55                let output = Output::Stream(m.content.clone());
56                if self.clear_on_next_output {
57                    self.clear_output().await;
58                    self.clear_on_next_output = false;
59                }
60                self.add_output(output).await;
61            }
62            Response::DisplayData(m) => {
63                let output = Output::DisplayData(m.content.clone());
64                if self.clear_on_next_output {
65                    self.clear_output().await;
66                    self.clear_on_next_output = false;
67                }
68                self.add_output(output).await;
69            }
70            Response::Error(m) => {
71                let output = Output::Error(m.content.clone());
72                if self.clear_on_next_output {
73                    self.clear_output().await;
74                    self.clear_on_next_output = false;
75                }
76                self.add_output(output).await;
77            }
78            Response::ClearOutput(m) => {
79                if m.content.wait {
80                    self.clear_on_next_output = true;
81                } else {
82                    self.clear_output().await;
83                }
84            }
85            _ => {}
86        }
87    }
88}
89
90// SimpleOutputHandler doesn't update a document model, just stores a list of outputs in memory.
91// Useful for tests and maybe debug / demos, probably not something you care about for app building
92#[derive(Debug, Clone)]
93pub struct SimpleOutputHandler {
94    // interior mutability here because .handle needs to set this and is &self, and when trying
95    // to change that to &mut self then it broke the delegation of ZMQ messages to Actions over
96    // in actions.rs. TODO: come back to this when I'm better at Rust?
97    clear_on_next_output: bool,
98    pub output: Vec<Output>,
99}
100
101impl Default for SimpleOutputHandler {
102    fn default() -> Self {
103        Self::new()
104    }
105}
106
107impl SimpleOutputHandler {
108    pub fn new() -> Self {
109        Self {
110            clear_on_next_output: false,
111            output: vec![],
112        }
113    }
114
115    async fn add_output(&mut self, content: Output) {
116        self.output.push(content);
117        println!("adding output");
118    }
119
120    async fn clear_output(&mut self) {
121        self.output.clear();
122        println!("clearing output");
123    }
124}
125
126#[async_trait::async_trait]
127impl Handler for SimpleOutputHandler {
128    async fn handle(&mut self, msg: &Response) {
129        match msg {
130            Response::ExecuteResult(m) => {
131                let output = Output::ExecuteResult(m.content.clone());
132                if self.clear_on_next_output {
133                    self.clear_output().await;
134                    self.clear_on_next_output = false;
135                }
136                self.add_output(output).await;
137            }
138            Response::Stream(m) => {
139                let output = Output::Stream(m.content.clone());
140                if self.clear_on_next_output {
141                    self.clear_output().await;
142                    self.clear_on_next_output = false;
143                }
144                self.add_output(output).await;
145            }
146            Response::DisplayData(m) => {
147                let output = Output::DisplayData(m.content.clone());
148                if self.clear_on_next_output {
149                    self.clear_output().await;
150                    self.clear_on_next_output = false;
151                }
152                self.add_output(output).await;
153            }
154            Response::Error(m) => {
155                let output = Output::Error(m.content.clone());
156                if self.clear_on_next_output {
157                    self.clear_output().await;
158                    self.clear_on_next_output = false;
159                }
160                self.add_output(output).await;
161            }
162            Response::ClearOutput(m) => {
163                if m.content.wait {
164                    self.clear_on_next_output = true;
165                } else {
166                    self.clear_output().await;
167                }
168            }
169            _ => {}
170        }
171    }
172}
kernel_sidecar/handlers/outputs.rs

kernel_sidecar/handlers/
outputs.rs
