Refactor Handlers

.github/workflows/ci.yaml

@@ -93,6 +93,14 @@ jobs:
     steps:
       - uses: actions/checkout@v3
 
+      # Deno needs Python/Jupyter in PATH before it can run `deno jupyter ...`
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+
+      - name: Install Jupyter
+        run: python -m pip install jupyter
+
       - name: Set up Deno stable release
         uses: denoland/setup-deno@v1
         with:

src/actions.rs

@@ -52,7 +52,7 @@ pub struct Action {
 impl Action {
     pub fn new(
         request: Request,
-        handlers: Vec<Arc<dyn Handler>>,
+        handlers: Vec<Arc<Mutex<dyn Handler>>>,
         msg_rx: mpsc::Receiver<Response>,
     ) -> Self {
         let action_state = Arc::new(Mutex::new(ActionState {
@@ -76,7 +76,7 @@ impl Action {
     async fn listen(
         mut msg_rx: mpsc::Receiver<Response>,
         expected_reply: ExpectedReplyType,
-        handlers: Vec<Arc<dyn Handler>>,
+        handlers: Vec<Arc<Mutex<dyn Handler>>>,
         action_state: Arc<Mutex<ActionState>>,
     ) {
         // We "finish" this background task when kernel idle and expected reply (if relevant) seen
@@ -87,7 +87,8 @@ impl Action {
             ExpectedReplyType::None => true,
         };
         while let Some(response) = msg_rx.recv().await {
-            for handler in &handlers {
+            for handler_arc in &handlers {
+                let mut handler = handler_arc.lock().await;
                 handler.handle(&response).await;
             }
             match response {

src/client.rs

@@ -39,7 +39,7 @@ use std::collections::HashMap;
 
 use std::sync::Arc;
 use std::time::Duration;
-use tokio::sync::{mpsc, Notify, RwLock};
+use tokio::sync::{mpsc, Mutex, Notify, RwLock};
 use tokio::time::sleep;
 use zeromq::{DealerSocket, ReqSocket, Socket, SocketRecv, SocketSend, SubSocket, ZmqMessage};
 
@@ -141,7 +141,11 @@ impl Client {
     // Creates an Action from a request + handlers, serializes the request to be sent over ZMQ,
     // sends over shell channel, and registers the request header msg_id in the Actions hashmap
     // so that all response messages can get routed to the appropriate Action handlers
-    async fn send_request(&self, request: Request, handlers: Vec<Arc<dyn Handler>>) -> Action {
+    async fn send_request(
+        &self,
+        request: Request,
+        handlers: Vec<Arc<Mutex<dyn Handler>>>,
+    ) -> Action {
         let (msg_tx, msg_rx) = mpsc::channel(100);
         let action = Action::new(request, handlers, msg_rx);
         let msg_id = action.request.msg_id();
@@ -152,12 +156,16 @@ impl Client {
         action
     }
 
-    pub async fn kernel_info_request(&self, handlers: Vec<Arc<dyn Handler>>) -> Action {
+    pub async fn kernel_info_request(&self, handlers: Vec<Arc<Mutex<dyn Handler>>>) -> Action {
         let request = KernelInfoRequest::new();
         self.send_request(request.into(), handlers).await
     }
 
-    pub async fn execute_request(&self, code: String, handlers: Vec<Arc<dyn Handler>>) -> Action {
+    pub async fn execute_request(
+        &self,
+        code: String,
+        handlers: Vec<Arc<Mutex<dyn Handler>>>,
+    ) -> Action {
         let request = ExecuteRequest::new(code);
         self.send_request(request.into(), handlers).await
     }

src/handlers/debug.rs

@@ -20,7 +20,7 @@ impl DebugHandler {
 
 #[async_trait::async_trait]
 impl Handler for DebugHandler {
-    async fn handle(&self, msg: &Response) {
+    async fn handle(&mut self, msg: &Response) {
         dbg!(msg);
     }
 }

src/handlers/mod.rs

@@ -12,5 +12,5 @@ pub use outputs::SimpleOutputHandler;
 
 #[async_trait::async_trait]
 pub trait Handler: Debug + Send + Sync {
-    async fn handle(&self, msg: &Response);
+    async fn handle(&mut self, msg: &Response);
 }

src/handlers/msg_count.rs

@@ -1,7 +1,4 @@
 use std::collections::HashMap;
-use std::sync::Arc;
-
-use tokio::sync::Mutex;
 
 use crate::jupyter::response::Response;
 
@@ -11,7 +8,7 @@ use crate::handlers::Handler;
 // Primarily used in tests and introspective click-testing
 #[derive(Debug, Clone)]
 pub struct MessageCountHandler {
-    pub counts: Arc<Mutex<HashMap<String, usize>>>,
+    pub counts: HashMap<String, usize>,
 }
 
 impl Default for MessageCountHandler {
@@ -23,17 +20,16 @@ impl Default for MessageCountHandler {
 impl MessageCountHandler {
     pub fn new() -> Self {
         MessageCountHandler {
-            counts: Arc::new(Mutex::new(HashMap::new())),
+            counts: HashMap::new(),
         }
     }
 }
 
 #[async_trait::async_trait]
 impl Handler for MessageCountHandler {
-    async fn handle(&self, msg: &Response) {
-        let mut counts = self.counts.lock().await;
+    async fn handle(&mut self, msg: &Response) {
         let msg_type = msg.msg_type();
-        let count = counts.entry(msg_type).or_insert(0);
+        let count = self.counts.entry(msg_type).or_insert(0);
         *count += 1;
     }
 }

src/handlers/outputs.rs

@@ -1,19 +1,101 @@
-use tokio::sync::{Mutex, RwLock};
+use tokio::sync::Mutex;
 
 use crate::handlers::Handler;
 use crate::jupyter::response::Response;
-use crate::notebook::Output;
+use crate::notebook::{Notebook, Output};
 
 use std::fmt::Debug;
 use std::sync::Arc;
 
+// Update a document model with outputs while running a cell
+#[derive(Debug)]
+pub struct OutputHandler {
+    nb: Arc<Mutex<Notebook>>,
+    cell_id: String,
+    clear_on_next_output: bool,
+}
+
+impl OutputHandler {
+    pub fn new(nb: Arc<Mutex<Notebook>>, cell_id: &str) -> Self {
+        Self {
+            nb,
+            cell_id: cell_id.to_string(),
+            clear_on_next_output: false,
+        }
+    }
+
+    pub async fn add_output(&mut self, content: Output) {
+        let mut nb = self.nb.lock().await;
+        if let Some(cell) = nb.get_mut_cell(&self.cell_id) {
+            cell.add_output(content);
+        }
+    }
+
+    pub async fn clear_output(&mut self) {
+        let mut nb = self.nb.lock().await;
+        if let Some(cell) = nb.get_mut_cell(&self.cell_id) {
+            cell.clear_output();
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl Handler for OutputHandler {
+    async fn handle(&mut self, msg: &Response) {
+        match msg {
+            Response::ExecuteResult(m) => {
+                let output = Output::ExecuteResult(m.content.clone());
+                if self.clear_on_next_output {
+                    self.clear_output().await;
+                    self.clear_on_next_output = false;
+                }
+                self.add_output(output).await;
+            }
+            Response::Stream(m) => {
+                let output = Output::Stream(m.content.clone());
+                if self.clear_on_next_output {
+                    self.clear_output().await;
+                    self.clear_on_next_output = false;
+                }
+                self.add_output(output).await;
+            }
+            Response::DisplayData(m) => {
+                let output = Output::DisplayData(m.content.clone());
+                if self.clear_on_next_output {
+                    self.clear_output().await;
+                    self.clear_on_next_output = false;
+                }
+                self.add_output(output).await;
+            }
+            Response::Error(m) => {
+                let output = Output::Error(m.content.clone());
+                if self.clear_on_next_output {
+                    self.clear_output().await;
+                    self.clear_on_next_output = false;
+                }
+                self.add_output(output).await;
+            }
+            Response::ClearOutput(m) => {
+                if m.content.wait {
+                    self.clear_on_next_output = true;
+                } else {
+                    self.clear_output().await;
+                }
+            }
+            _ => {}
+        }
+    }
+}
+
+// SimpleOutputHandler doesn't update a document model, just stores a list of outputs in memory.
+// Useful for tests and maybe debug / demos, probably not something you care about for app building
 #[derive(Debug, Clone)]
 pub struct SimpleOutputHandler {
     // interior mutability here because .handle needs to set this and is &self, and when trying
     // to change that to &mut self then it broke the delegation of ZMQ messages to Actions over
     // in actions.rs. TODO: come back to this when I'm better at Rust?
-    clear_on_next_output: Arc<Mutex<bool>>,
-    pub output: Arc<RwLock<Vec<Output>>>,
+    clear_on_next_output: bool,
+    pub output: Vec<Output>,
 }
 
 impl Default for SimpleOutputHandler {
@@ -25,64 +107,63 @@ impl Default for SimpleOutputHandler {
 impl SimpleOutputHandler {
     pub fn new() -> Self {
         Self {
-            clear_on_next_output: Arc::new(Mutex::new(false)),
-            output: Arc::new(RwLock::new(vec![])),
+            clear_on_next_output: false,
+            output: vec![],
         }
     }
 
-    async fn add_cell_content(&self, content: Output) {
-        self.output.write().await.push(content);
-        println!("add_cell_content");
+    async fn add_output(&mut self, content: Output) {
+        self.output.push(content);
+        println!("adding output");
     }
 
-    async fn clear_cell_content(&self) {
-        self.output.write().await.clear();
-        println!("clear_cell_content");
+    async fn clear_output(&mut self) {
+        self.output.clear();
+        println!("clearing output");
     }
 }
 
 #[async_trait::async_trait]
 impl Handler for SimpleOutputHandler {
-    async fn handle(&self, msg: &Response) {
-        let mut clear_on_next_output = self.clear_on_next_output.lock().await;
+    async fn handle(&mut self, msg: &Response) {
         match msg {
             Response::ExecuteResult(m) => {
                 let output = Output::ExecuteResult(m.content.clone());
-                if *clear_on_next_output {
-                    self.clear_cell_content().await;
-                    *clear_on_next_output = false;
+                if self.clear_on_next_output {
+                    self.clear_output().await;
+                    self.clear_on_next_output = false;
                 }
-                self.add_cell_content(output).await;
+                self.add_output(output).await;
             }
             Response::Stream(m) => {
                 let output = Output::Stream(m.content.clone());
-                if *clear_on_next_output {
-                    self.clear_cell_content().await;
-                    *clear_on_next_output = false;
+                if self.clear_on_next_output {
+                    self.clear_output().await;
+                    self.clear_on_next_output = false;
                 }
-                self.add_cell_content(output).await;
+                self.add_output(output).await;
             }
             Response::DisplayData(m) => {
                 let output = Output::DisplayData(m.content.clone());
-                if *clear_on_next_output {
-                    self.clear_cell_content().await;
-                    *clear_on_next_output = false;
+                if self.clear_on_next_output {
+                    self.clear_output().await;
+                    self.clear_on_next_output = false;
                 }
-                self.add_cell_content(output).await;
+                self.add_output(output).await;
             }
             Response::Error(m) => {
                 let output = Output::Error(m.content.clone());
-                if *clear_on_next_output {
-                    self.clear_cell_content().await;
-                    *clear_on_next_output = false;
+                if self.clear_on_next_output {
+                    self.clear_output().await;
+                    self.clear_on_next_output = false;
                 }
-                self.add_cell_content(output).await;
+                self.add_output(output).await;
             }
             Response::ClearOutput(m) => {
                 if m.content.wait {
-                    *clear_on_next_output = true;
+                    self.clear_on_next_output = true;
                 } else {
-                    self.clear_cell_content().await;
+                    self.clear_output().await;
                 }
             }
             _ => {}

src/lib.rs

@@ -3,5 +3,4 @@ pub mod client;
 pub mod handlers;
 pub mod jupyter;
 pub mod kernels;
-pub mod nb_builder;
 pub mod notebook;