1use std::io::{self, BufWriter, Stdout, Write};
2
3use crossterm::{
4 cursor,
5 event::{
6 DisableBracketedPaste, DisableFocusChange, DisableMouseCapture, EnableBracketedPaste,
7 EnableFocusChange, EnableMouseCapture, KeyboardEnhancementFlags,
8 PopKeyboardEnhancementFlags, PushKeyboardEnhancementFlags,
9 },
10 terminal::{self, DisableLineWrap, EnableLineWrap, EnterAlternateScreen, LeaveAlternateScreen},
11 QueueableCommand,
12};
13
14#[derive(Clone, Copy, Debug, PartialEq, Eq)]
15pub enum SuspendScreen {
16 KeepAlternate,
17 LeaveAlternate,
18}
19
20pub struct TerminalSession<W: Write> {
21 writer: W,
22 raw_mode: bool,
23 modes: TerminalModes,
24 _not_send: std::marker::PhantomData<*const ()>,
25}
26
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28struct TerminalModes {
29 alternate_screen: bool,
30 mouse_capture: bool,
31 hide_cursor: bool,
32 keyboard_enhancements: Option<KeyboardEnhancementFlags>,
33 bracketed_paste: bool,
34 focus_events: bool,
35 line_wrap: bool,
36}
37
38impl TerminalModes {
39 fn inactive() -> Self {
40 Self {
41 alternate_screen: false,
42 mouse_capture: false,
43 hide_cursor: false,
44 keyboard_enhancements: None,
45 bracketed_paste: false,
46 focus_events: false,
47 line_wrap: true,
48 }
49 }
50}
51
52impl Default for TerminalModes {
53 fn default() -> Self {
54 Self {
55 alternate_screen: true,
56 mouse_capture: true,
57 hide_cursor: true,
58 keyboard_enhancements: None,
59 bracketed_paste: false,
60 focus_events: false,
61 line_wrap: true,
62 }
63 }
64}
65
66#[derive(Clone, Debug)]
67pub struct TerminalSessionBuilder {
68 modes: TerminalModes,
69 buffer_capacity: usize,
70}
71
72impl Default for TerminalSessionBuilder {
73 fn default() -> Self {
74 Self {
75 modes: TerminalModes::default(),
76 buffer_capacity: 64 * 1024,
77 }
78 }
79}
80
81impl TerminalSessionBuilder {
82 pub fn alternate_screen(mut self, yes: bool) -> Self {
83 self.modes.alternate_screen = yes;
84 self
85 }
86
87 pub fn mouse_capture(mut self, yes: bool) -> Self {
88 self.modes.mouse_capture = yes;
89 self
90 }
91
92 pub fn hide_cursor(mut self, yes: bool) -> Self {
93 self.modes.hide_cursor = yes;
94 self
95 }
96
97 pub fn keyboard_enhancements(mut self, flags: KeyboardEnhancementFlags) -> Self {
98 self.modes.keyboard_enhancements = Some(flags);
99 self
100 }
101
102 pub fn bracketed_paste(mut self, yes: bool) -> Self {
103 self.modes.bracketed_paste = yes;
104 self
105 }
106
107 pub fn focus_events(mut self, yes: bool) -> Self {
108 self.modes.focus_events = yes;
109 self
110 }
111
112 pub fn line_wrap(mut self, yes: bool) -> Self {
113 self.modes.line_wrap = yes;
114 self
115 }
116
117 pub fn buffer_capacity(mut self, bytes: usize) -> Self {
118 self.buffer_capacity = bytes;
119 self
120 }
121
122 pub fn enter_stdout(self) -> io::Result<TerminalSession<BufWriter<Stdout>>> {
123 let capacity = self.buffer_capacity;
124 let stdout = io::stdout();
125 self.enter(BufWriter::with_capacity(capacity, stdout))
126 }
127
128 pub fn enter<W: Write>(self, writer: W) -> io::Result<TerminalSession<W>> {
129 TerminalSession::enter_with(writer, self)
130 }
131}
132
133impl TerminalSession<BufWriter<Stdout>> {
134 pub fn enter_stdout() -> io::Result<Self> {
135 TerminalSession::builder().enter_stdout()
136 }
137
138 pub fn builder() -> TerminalSessionBuilder {
139 TerminalSessionBuilder::default()
140 }
141}
142
143impl<W: Write> TerminalSession<W> {
144 fn enter_with(writer: W, options: TerminalSessionBuilder) -> io::Result<Self> {
145 terminal::enable_raw_mode()?;
146 let mut session = Self {
147 writer,
148 raw_mode: true,
149 modes: TerminalModes::inactive(),
150 _not_send: std::marker::PhantomData,
151 };
152
153 if let Err(err) = session.enter_modes(options.modes) {
154 let _ = session.restore();
155 return Err(err);
156 }
157
158 Ok(session)
159 }
160
161 pub fn writer(&mut self) -> &mut W {
162 &mut self.writer
163 }
164
165 pub fn size(&self) -> io::Result<(u16, u16)> {
166 terminal::size()
167 }
168
169 pub fn suspend<F, R>(&mut self, f: F) -> R
170 where
171 F: FnOnce() -> R,
172 {
173 self.suspend_with(SuspendScreen::KeepAlternate, f)
174 }
175
176 pub fn suspend_with<F, R>(&mut self, screen: SuspendScreen, f: F) -> R
177 where
178 F: FnOnce() -> R,
179 {
180 let modes = self.modes;
181 let raw_mode = self.raw_mode;
182 let _ = self.release_input_modes();
183 if screen == SuspendScreen::LeaveAlternate && modes.alternate_screen {
184 let _ = self.writer.queue(LeaveAlternateScreen);
185 let _ = self.writer.flush();
186 }
187 let result = f();
188 if modes.alternate_screen {
189 let _ = self.writer.queue(EnterAlternateScreen);
190 }
191 self.modes = modes;
192 let _ = self.restore_input_modes(raw_mode);
193 result
194 }
195
196 fn enter_modes(&mut self, modes: TerminalModes) -> io::Result<()> {
197 if modes.alternate_screen {
198 self.writer.queue(EnterAlternateScreen)?;
199 self.modes.alternate_screen = true;
200 }
201 if !modes.line_wrap {
202 self.writer.queue(DisableLineWrap)?;
203 self.modes.line_wrap = false;
204 }
205 if modes.hide_cursor {
206 self.writer.queue(cursor::Hide)?;
207 self.modes.hide_cursor = true;
208 }
209 if modes.bracketed_paste {
210 self.writer.queue(EnableBracketedPaste)?;
211 self.modes.bracketed_paste = true;
212 }
213 if modes.focus_events {
214 self.writer.queue(EnableFocusChange)?;
215 self.modes.focus_events = true;
216 }
217 if let Some(flags) = modes.keyboard_enhancements {
218 self.writer.queue(PushKeyboardEnhancementFlags(flags))?;
219 self.modes.keyboard_enhancements = Some(flags);
220 }
221 if modes.mouse_capture {
222 self.writer.queue(EnableMouseCapture)?;
223 self.modes.mouse_capture = true;
224 }
225 self.writer.flush()
226 }
227
228 fn restore(&mut self) -> io::Result<()> {
229 self.release_input_modes()?;
230 if self.modes.alternate_screen {
231 self.writer.queue(LeaveAlternateScreen)?;
232 self.modes.alternate_screen = false;
233 }
234 self.writer.flush()?;
235 if self.raw_mode {
236 terminal::disable_raw_mode()?;
237 self.raw_mode = false;
238 }
239 Ok(())
240 }
241
242 fn release_input_modes(&mut self) -> io::Result<()> {
243 if self.modes.mouse_capture {
244 self.writer.queue(DisableMouseCapture)?;
245 self.modes.mouse_capture = false;
246 }
247 if self.modes.keyboard_enhancements.is_some() {
248 self.writer.queue(PopKeyboardEnhancementFlags)?;
249 self.modes.keyboard_enhancements = None;
250 }
251 if self.modes.focus_events {
252 self.writer.queue(DisableFocusChange)?;
253 self.modes.focus_events = false;
254 }
255 if self.modes.bracketed_paste {
256 self.writer.queue(DisableBracketedPaste)?;
257 self.modes.bracketed_paste = false;
258 }
259 if self.modes.hide_cursor {
260 self.writer.queue(cursor::Show)?;
261 self.modes.hide_cursor = false;
262 }
263 if !self.modes.line_wrap {
264 self.writer.queue(EnableLineWrap)?;
265 self.modes.line_wrap = true;
266 }
267 self.writer.flush()?;
268 if self.raw_mode {
269 terminal::disable_raw_mode()?;
270 self.raw_mode = false;
271 }
272 Ok(())
273 }
274
275 fn restore_input_modes(&mut self, raw_mode: bool) -> io::Result<()> {
276 if raw_mode {
277 terminal::enable_raw_mode()?;
278 self.raw_mode = true;
279 }
280 if !self.modes.line_wrap {
281 self.writer.queue(DisableLineWrap)?;
282 }
283 if self.modes.hide_cursor {
284 self.writer.queue(cursor::Hide)?;
285 }
286 if self.modes.bracketed_paste {
287 self.writer.queue(EnableBracketedPaste)?;
288 }
289 if self.modes.focus_events {
290 self.writer.queue(EnableFocusChange)?;
291 }
292 if let Some(flags) = self.modes.keyboard_enhancements {
293 self.writer.queue(PushKeyboardEnhancementFlags(flags))?;
294 }
295 if self.modes.mouse_capture {
296 self.writer.queue(EnableMouseCapture)?;
297 }
298 self.writer.flush()
299 }
300}
301
302impl<W: Write> Drop for TerminalSession<W> {
303 fn drop(&mut self) {
304 let _ = self.restore();
305 }
306}
307
308#[cfg(test)]
309mod tests {
310 use super::*;
311
312 #[derive(Default)]
313 struct TestWriter {
314 bytes: Vec<u8>,
315 fail_flush: bool,
316 }
317
318 impl Write for TestWriter {
319 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
320 self.bytes.extend_from_slice(buf);
321 Ok(buf.len())
322 }
323
324 fn flush(&mut self) -> io::Result<()> {
325 if self.fail_flush {
326 Err(io::Error::other("flush failed"))
327 } else {
328 Ok(())
329 }
330 }
331 }
332
333 fn test_session(writer: TestWriter) -> TerminalSession<TestWriter> {
334 TerminalSession {
335 writer,
336 raw_mode: false,
337 modes: TerminalModes::inactive(),
338 _not_send: std::marker::PhantomData,
339 }
340 }
341
342 fn all_modes() -> TerminalModes {
343 TerminalModes {
344 alternate_screen: true,
345 mouse_capture: true,
346 hide_cursor: true,
347 keyboard_enhancements: Some(KeyboardEnhancementFlags::DISAMBIGUATE_ESCAPE_CODES),
348 bracketed_paste: true,
349 focus_events: true,
350 line_wrap: false,
351 }
352 }
353
354 #[test]
355 fn terminal_session_tracks_and_restores_modes() {
356 let mut session = test_session(TestWriter::default());
357 session.enter_modes(all_modes()).unwrap();
358 assert_eq!(session.modes, all_modes());
359 assert!(!session.writer.bytes.is_empty());
360
361 session.restore().unwrap();
362 assert_eq!(session.modes, TerminalModes::inactive());
363 }
364
365 #[test]
366 fn terminal_session_can_restore_after_enter_flush_failure() {
367 let mut session = test_session(TestWriter {
368 fail_flush: true,
369 ..TestWriter::default()
370 });
371
372 let err = session.enter_modes(all_modes()).unwrap_err();
373 assert_eq!(err.kind(), io::ErrorKind::Other);
374 assert_eq!(session.modes, all_modes());
375
376 session.writer.fail_flush = false;
377 session.restore().unwrap();
378 assert_eq!(session.modes, TerminalModes::inactive());
379 }
380
381 #[test]
382 fn terminal_session_suspend_can_leave_alternate_screen() {
383 let mut session = test_session(TestWriter::default());
384 session.modes = all_modes();
385
386 let result = session.suspend_with(SuspendScreen::LeaveAlternate, || 42);
387
388 assert_eq!(result, 42);
389 assert_eq!(session.modes, all_modes());
390 let out = String::from_utf8_lossy(&session.writer.bytes);
391 assert!(out.contains("1049l"), "missing leave alt screen: {out:?}");
392 assert!(out.contains("1049h"), "missing enter alt screen: {out:?}");
393 }
394
395 #[test]
396 fn terminal_session_suspend_keeps_alternate_screen_by_default() {
397 let mut session = test_session(TestWriter::default());
398 session.modes = all_modes();
399
400 session.suspend(|| ());
401
402 let out = String::from_utf8_lossy(&session.writer.bytes);
403 assert!(
404 !out.contains("1049l"),
405 "unexpected leave alt screen: {out:?}"
406 );
407 assert!(out.contains("1049h"), "missing enter alt screen: {out:?}");
408 }
409}