Add Terminal::set_panic_hook

src/lib.rs

@@ -9,4 +9,4 @@ pub use event::{
     stream::{DummyEventStream, EventStream},
     Event,
 };
-pub use terminal::{PlatformTerminal, Terminal};
+pub use terminal::{PlatformHandle, PlatformTerminal, Terminal};

src/terminal.rs

@@ -22,6 +22,11 @@ pub type PlatformTerminal = UnixTerminal;
 #[cfg(windows)]
 pub type PlatformTerminal = WindowsTerminal;
 
+#[cfg(unix)]
+pub type PlatformHandle = FileDescriptor;
+#[cfg(windows)]
+pub type PlatformHandle = OutputHandle;
+
 // CREDIT: This is heavily based on termwiz.
 // <https://github.com/wezterm/wezterm/blob/a87358516004a652ad840bc1661bdf65ffc89b43/termwiz/src/terminal/mod.rs#L50-L111>
 // This trait is simpler, however, and the terminals themselves do not have drop glue or try
@@ -59,4 +64,12 @@ pub trait Terminal: io::Write {
     /// This function blocks until an `Event` is available. Use `poll` first to guarantee that the
     /// read won't block.
     fn read<F: Fn(&Event) -> bool>(&self, filter: F) -> io::Result<Event>;
+    /// Sets a hook function to run.
+    ///
+    /// Depending on how your application handles panics you may wish to set a panic hook which
+    /// eagerly resets the terminal (such as by disabling bracketed paste and entering the main
+    /// screen). The parameter for this hook is a platform handle to `std::io::stdout` or
+    /// equivalent which implements `std::io::Write`. When the hook function is finished running
+    /// the handle's modes will be reset (same as `enter_cooked_mode`).
+    fn set_panic_hook(&mut self, f: impl Fn(&mut PlatformHandle) + Send + Sync + 'static);
 }

src/terminal/unix.rs

@@ -19,7 +19,7 @@ const BUF_SIZE: usize = 4096;
 // <https://github.com/wezterm/wezterm/blob/a87358516004a652ad840bc1661bdf65ffc89b43/filedescriptor/src/unix.rs>
 
 #[derive(Debug)]
-pub(crate) enum FileDescriptor {
+pub enum FileDescriptor {
     Owned(OwnedFd),
     Borrowed(BorrowedFd<'static>),
 }
@@ -101,6 +101,7 @@ pub struct UnixTerminal {
     write: BufWriter<FileDescriptor>,
     /// The termios of the PTY's writer detected during `Self::new`.
     original_termios: Termios,
+    has_panic_hook: bool,
 }
 
 impl UnixTerminal {
@@ -114,6 +115,7 @@ impl UnixTerminal {
             reader,
             write: BufWriter::with_capacity(BUF_SIZE, write),
             original_termios,
+            has_panic_hook: false,
         })
     }
 }
@@ -163,12 +165,27 @@ impl Terminal for UnixTerminal {
     fn read<F: Fn(&Event) -> bool>(&self, filter: F) -> io::Result<Event> {
         self.reader.read(filter)
     }
+
+    fn set_panic_hook(&mut self, f: impl Fn(&mut FileDescriptor) + Send + Sync + 'static) {
+        let original_termios = self.original_termios.clone();
+        let hook = std::panic::take_hook();
+        std::panic::set_hook(Box::new(move |info| {
+            if let Ok((_read, mut write)) = open_pty() {
+                f(&mut write);
+                let _ = termios::tcsetattr(write, termios::OptionalActions::Now, &original_termios);
+            }
+            hook(info);
+        }));
+        self.has_panic_hook = true;
+    }
 }
 
 impl Drop for UnixTerminal {
     fn drop(&mut self) {
-        let _ = self.flush();
-        let _ = self.enter_cooked_mode();
+        if !self.has_panic_hook || !std::thread::panicking() {
+            let _ = self.flush();
+            let _ = self.enter_cooked_mode();
+        }
     }
 }
 

src/terminal/windows.rs

@@ -153,7 +153,7 @@ impl AsRawHandle for InputHandle {
 }
 
 #[derive(Debug)]
-pub(crate) struct OutputHandle {
+pub struct OutputHandle {
     handle: OwnedHandle,
 }
 
@@ -281,6 +281,7 @@ pub struct WindowsTerminal {
     original_output_mode: CONSOLE_MODE,
     original_input_cp: CodePageID,
     original_output_cp: CodePageID,
+    has_panic_hook: bool,
 }
 
 impl WindowsTerminal {
@@ -317,6 +318,7 @@ impl WindowsTerminal {
             original_output_mode,
             original_input_cp,
             original_output_cp,
+            has_panic_hook: false,
         })
     }
 }
@@ -382,15 +384,36 @@ impl Terminal for WindowsTerminal {
     fn read<F: Fn(&Event) -> bool>(&self, filter: F) -> io::Result<Event> {
         self.reader.read(filter)
     }
+
+    fn set_panic_hook(&mut self, f: impl Fn(&mut OutputHandle) + Send + Sync + 'static) {
+        let original_input_cp = self.original_input_cp;
+        let original_input_mode = self.original_input_mode;
+        let original_output_cp = self.original_output_cp;
+        let original_output_mode = self.original_output_mode;
+        let hook = std::panic::take_hook();
+        std::panic::set_hook(Box::new(move |info| {
+            if let Ok((mut input, mut output)) = open_pty() {
+                f(&mut output);
+                let _ = input.set_code_page(original_input_cp);
+                let _ = input.set_mode(original_input_mode);
+                let _ = output.set_code_page(original_output_cp);
+                let _ = output.set_mode(original_output_mode);
+            }
+            hook(info);
+        }));
+        self.has_panic_hook = true;
+    }
 }
 
 impl Drop for WindowsTerminal {
     fn drop(&mut self) {
-        let _ = self.flush();
-        let _ = self.input.set_code_page(self.original_input_cp);
-        let _ = self.output.get_mut().set_code_page(self.original_output_cp);
-        let _ = self.input.set_mode(self.original_input_mode);
-        let _ = self.output.get_mut().set_mode(self.original_output_mode);
+        if !self.has_panic_hook || !std::thread::panicking() {
+            let _ = self.flush();
+            let _ = self.input.set_code_page(self.original_input_cp);
+            let _ = self.output.get_mut().set_code_page(self.original_output_cp);
+            let _ = self.input.set_mode(self.original_input_mode);
+            let _ = self.output.get_mut().set_mode(self.original_output_mode);
+        }
     }
 }