async_tty/
terminal.rs

1use std::io::{Stdin, Stdout, stdin, stdout};
2use std::os::fd::AsFd;
3use std::task::Poll;
4
5use rustix::io::write;
6use rustix::termios::{OptionalActions, Termios, tcgetattr, tcsetattr};
7use snafu::ResultExt;
8use tokio::io::{AsyncWrite, AsyncWriteExt};
9
10use crate::Result;
11use crate::error::{
12    GetTerminalAttributesSnafu, SetTerminalAttributesSnafu, SwitchToAlternateScreenSnafu,
13    SwitchToMainScreenSnafu,
14};
15
16pub struct Terminal<In: AsFd, Out: AsFd> {
17    pub(crate) stdin: In,
18    pub(crate) stdout: Out,
19    pub(crate) original_termios: Termios,
20}
21
22// Implement Unpin to allow the struct to be safely unpinned
23impl<In: AsFd, Out: AsFd> Unpin for Terminal<In, Out> {}
24
25impl Terminal<Stdin, Stdout> {
26    pub fn new() -> Result<Self> {
27        let stdin = stdin();
28        let stdout = stdout();
29        let original_termios = tcgetattr(&stdin).context(GetTerminalAttributesSnafu)?;
30        Ok(Terminal {
31            stdin,
32            stdout,
33            original_termios,
34        })
35    }
36}
37
38impl<In: AsFd, Out: AsFd> Terminal<In, Out> {
39    /// Set the terminal to raw mode
40    pub fn set_raw_mode(&self) -> Result<()> {
41        let mut raw = self.original_termios.clone();
42        raw.make_raw();
43        // TODO work out whether we should be using Drain or Flush here instead of Now
44        tcsetattr(&self.stdin, OptionalActions::Now, &raw).context(SetTerminalAttributesSnafu)?;
45        Ok(())
46    }
47
48    /// Set the terminal to cooked mode
49    pub fn set_cooked_mode(&self) -> Result<()> {
50        // TODO work out whether we should be using Drain or Flush here instead of Now
51        tcsetattr(&self.stdin, OptionalActions::Now, &self.original_termios)
52            .context(SetTerminalAttributesSnafu)?;
53        Ok(())
54    }
55
56    /// Switch to alternate screen buffer
57    pub async fn switch_to_alternate_screen(&mut self) -> Result<()> {
58        self.write(b"\x1b[?1049h")
59            .await
60            .context(SwitchToAlternateScreenSnafu)?;
61        Ok(())
62    }
63
64    /// Switch back to main screen buffer
65    pub async fn switch_to_main_screen(&mut self) -> Result<()> {
66        self.write(b"\x1b[?1049l")
67            .await
68            .context(SwitchToMainScreenSnafu)?;
69        Ok(())
70    }
71}
72
73impl<In: AsFd, Out: AsFd> Drop for Terminal<In, Out> {
74    fn drop(&mut self) {
75        // Restore the original terminal attributes when the Terminal struct is dropped
76        if let Err(e) = self.set_cooked_mode() {
77            eprintln!("Failed to restore terminal attributes: {}", e);
78        }
79    }
80}
81
82impl<In: AsFd, Out: AsFd> AsyncWrite for Terminal<In, Out> {
83    fn poll_write(
84        self: std::pin::Pin<&mut Self>,
85        _cx: &mut std::task::Context<'_>,
86        buf: &[u8],
87    ) -> Poll<std::io::Result<usize>> {
88        let this = self.get_mut();
89        let result = write(&this.stdout, buf);
90        match result {
91            Ok(n) => Poll::Ready(Ok(n)),
92            Err(e) => Poll::Ready(Err(std::io::Error::other(e))),
93        }
94    }
95
96    fn poll_flush(
97        self: std::pin::Pin<&mut Self>,
98        _cx: &mut std::task::Context<'_>,
99    ) -> Poll<std::io::Result<()>> {
100        // Implement flush if needed
101        Poll::Ready(Ok(()))
102    }
103
104    fn poll_shutdown(
105        self: std::pin::Pin<&mut Self>,
106        _cx: &mut std::task::Context<'_>,
107    ) -> Poll<std::io::Result<()>> {
108        // Implement shutdown if needed
109        Poll::Ready(Ok(()))
110    }
111}