From 6f84f5f594e3e8e920d76da34f3293c8254ce611 Mon Sep 17 00:00:00 2001 From: Ling Wang Date: Wed, 18 Sep 2024 22:32:23 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat:=20Add=20deansi=20tty?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pythonapi/deansi.rs | 36 ++++++ src/term/deansi.rs | 53 ++++++++ src/vendor/mod.rs | 1 + src/vendor/strip_ansi_escapes.rs | 215 +++++++++++++++++++++++++++++++ 4 files changed, 305 insertions(+) create mode 100644 src/pythonapi/deansi.rs create mode 100644 src/term/deansi.rs create mode 100644 src/vendor/mod.rs create mode 100644 src/vendor/strip_ansi_escapes.rs diff --git a/src/pythonapi/deansi.rs b/src/pythonapi/deansi.rs new file mode 100644 index 0000000..32ff3c3 --- /dev/null +++ b/src/pythonapi/deansi.rs @@ -0,0 +1,36 @@ +use pyo3::{exceptions::PyRuntimeError, pyclass, pymethods, PyResult}; + +use crate::{term::tty::Tty, util::anybase::heap_raw}; + +use super::shell_like::{handle_wrap, PyTty, PyTtyWrapper, TtyType}; + +pub fn handle_deansi(inner: &mut Option) -> PyResult<()> { + if inner.is_none() { + return Err(PyRuntimeError::new_err( + "You must define at least one valid object", + )); + } + let mut be_wrapped = inner.take().unwrap(); + let be_wrapped = be_wrapped.safe_take()?; + let be_wrapped = Box::into_inner(be_wrapped); + let dean = Box::new(crate::term::deansi::DeANSI::build(be_wrapped)); + let dean: Box = dean as TtyType; + *inner = Some(PyTtyWrapper { + tty: heap_raw(dean), + }); + Ok(()) +} + +#[pyclass(extends=PyTty, subclass)] +pub struct DeANSI {} + +#[pymethods] +impl DeANSI { + #[new] + fn py_new(be_wrapped: &mut PyTty) -> PyResult<(Self, PyTty)> { + let mut inner = None; + handle_wrap(&mut inner, Some(be_wrapped))?; + handle_deansi(&mut inner)?; + Ok((DeANSI {}, PyTty::build(inner.unwrap()))) + } +} diff --git a/src/term/deansi.rs b/src/term/deansi.rs new file mode 100644 index 0000000..bb54c14 --- /dev/null +++ b/src/term/deansi.rs @@ -0,0 +1,53 @@ +use std::{any::Any, error::Error}; + +use crate::{util::anybase::AnyBase, vendor::strip_ansi_escapes}; + +use super::tty::{DynTty, Tty, WrapperTty}; + +pub struct DeANSI { + inner: DynTty, +} + +impl DeANSI { + pub fn build(inner: DynTty) -> DeANSI { + DeANSI { inner } + } +} + +impl AnyBase for DeANSI { + fn as_any(&self) -> &dyn Any { + self + } + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + fn into_any(self: Box) -> Box { + self + } +} + +impl Tty for DeANSI { + fn read(&mut self) -> Result, Box> { + // Due to the escape sequences may be cut off in the middle of the buffer, + // read the buffer in line is needed. + let data = self.inner.read_line()?; + let data = strip_ansi_escapes::strip(&data); + Ok(data) + } + fn read_line(&mut self) -> Result, Box> { + let data = self.inner.read_line()?; + let data = strip_ansi_escapes::strip(&data); + Ok(data) + } + fn write(&mut self, data: &[u8]) -> Result<(), Box> { + let data = strip_ansi_escapes::strip(&data); + self.inner.write(&data)?; + Ok(()) + } +} + +impl WrapperTty for DeANSI { + fn exit(self) -> DynTty { + self.inner + } +} diff --git a/src/vendor/mod.rs b/src/vendor/mod.rs new file mode 100644 index 0000000..300337e --- /dev/null +++ b/src/vendor/mod.rs @@ -0,0 +1 @@ +pub mod strip_ansi_escapes; \ No newline at end of file diff --git a/src/vendor/strip_ansi_escapes.rs b/src/vendor/strip_ansi_escapes.rs new file mode 100644 index 0000000..bbbfcc9 --- /dev/null +++ b/src/vendor/strip_ansi_escapes.rs @@ -0,0 +1,215 @@ +//! Originally from https://github.com/luser/strip-ansi-escapes +//! License: MIT/Apache-2.0 +//! A crate for stripping ANSI escape sequences from byte sequences. +//! +//! This can be used to take output from a program that includes escape sequences and write +//! it somewhere that does not easily support them, such as a log file. +//! +//! The simplest interface provided is the [`strip`] function, which takes a byte slice and returns +//! a `Vec` of bytes with escape sequences removed. For writing bytes directly to a writer, you +//! may prefer using the [`Writer`] struct, which implements `Write` and strips escape sequences +//! as they are written. +//! +//! [`strip`]: fn.strip.html +//! [`Writer`]: struct.Writer.html +//! +//! # Example +//! +//! ``` +//! use std::io::{self, Write}; +//! +//! # fn foo() -> io::Result<()> { +//! let bytes_with_colors = b"\x1b[32mfoo\x1b[m bar"; +//! let plain_bytes = strip_ansi_escapes::strip(&bytes_with_colors); +//! io::stdout().write_all(&plain_bytes)?; +//! # Ok(()) +//! # } +//! ``` + +use std::io::{self, Cursor, IntoInnerError, LineWriter, Write}; +use vte::{Parser, Perform}; + +/// `Writer` wraps an underlying type that implements `Write`, stripping ANSI escape sequences +/// from bytes written to it before passing them to the underlying writer. +/// +/// # Example +/// ``` +/// use std::io::{self, Write}; +/// use strip_ansi_escapes::Writer; +/// +/// # fn foo() -> io::Result<()> { +/// let bytes_with_colors = b"\x1b[32mfoo\x1b[m bar"; +/// let mut writer = Writer::new(io::stdout()); +/// // Only `foo bar` will be written to stdout +/// writer.write_all(bytes_with_colors)?; +/// # Ok(()) +/// # } +/// ``` + +pub struct Writer +where + W: Write, +{ + performer: Performer, + parser: Parser, +} + +/// Strip ANSI escapes from `data` and return the remaining bytes as a `Vec`. +/// +/// See [the module documentation][mod] for an example. +/// +/// [mod]: index.html +pub fn strip(data: T) -> Vec +where + T: AsRef<[u8]>, +{ + fn strip_impl(data: &[u8]) -> io::Result> { + let c = Cursor::new(Vec::new()); + let mut writer = Writer::new(c); + writer.write_all(data.as_ref())?; + Ok(writer.into_inner()?.into_inner()) + } + + strip_impl(data.as_ref()).expect("writing to a Cursor> cannot fail") +} + +/// Strip ANSI escapes from `data` and return the remaining contents as a `String`. +/// +/// # Example +/// +/// ``` +/// let str_with_colors = "\x1b[32mfoo\x1b[m bar"; +/// let string_without_colors = strip_ansi_escapes::strip_str(str_with_colors); +/// assert_eq!(string_without_colors, "foo bar"); +/// ``` +pub fn strip_str(data: T) -> String +where + T: AsRef, +{ + let bytes = strip(data.as_ref()); + String::from_utf8(bytes) + .expect("stripping ANSI escapes from a UTF-8 string always results in UTF-8") +} + +struct Performer +where + W: Write, +{ + writer: LineWriter, + err: Option, +} + +impl Writer +where + W: Write, +{ + /// Create a new `Writer` that writes to `inner`. + pub fn new(inner: W) -> Writer { + Writer { + performer: Performer { + writer: LineWriter::new(inner), + err: None, + }, + parser: Parser::new(), + } + } + + /// Unwraps this `Writer`, returning the underlying writer. + /// + /// The internal buffer is written out before returning the writer, which + /// may produce an [`IntoInnerError`]. + /// + /// [IntoInnerError]: https://doc.rust-lang.org/std/io/struct.IntoInnerError.html + pub fn into_inner(self) -> Result>> { + self.performer.into_inner() + } +} + +impl Write for Writer +where + W: Write, +{ + fn write(&mut self, buf: &[u8]) -> io::Result { + for b in buf.iter() { + self.parser.advance(&mut self.performer, *b) + } + match self.performer.err.take() { + Some(e) => Err(e), + None => Ok(buf.len()), + } + } + + fn flush(&mut self) -> io::Result<()> { + self.performer.flush() + } +} + +impl Performer +where + W: Write, +{ + pub fn flush(&mut self) -> io::Result<()> { + self.writer.flush() + } + + pub fn into_inner(self) -> Result>> { + self.writer.into_inner() + } +} + +impl Perform for Performer +where + W: Write, +{ + fn print(&mut self, c: char) { + // Just print bytes to the inner writer. + self.err = write!(self.writer, "{}", c).err(); + } + fn execute(&mut self, byte: u8) { + // We only care about executing linefeeds. + if byte == b'\n' { + self.err = writeln!(self.writer).err(); + } + } +} + +#[cfg(doctest)] +extern crate doc_comment; + +#[cfg(doctest)] +doc_comment::doctest!("../README.md", readme); + +#[cfg(test)] +mod tests { + use super::*; + + fn assert_parsed(input: &[u8], expected: &[u8]) { + let bytes = strip(input); + assert_eq!(bytes, expected); + } + + #[test] + fn test_simple() { + assert_parsed(b"\x1b[m\x1b[m\x1b[32m\x1b[1m Finished\x1b[m dev [unoptimized + debuginfo] target(s) in 0.0 secs", + b" Finished dev [unoptimized + debuginfo] target(s) in 0.0 secs"); + } + + #[test] + fn test_newlines() { + assert_parsed(b"foo\nbar\n", b"foo\nbar\n"); + } + + #[test] + fn test_escapes_newlines() { + assert_parsed(b"\x1b[m\x1b[m\x1b[32m\x1b[1m Compiling\x1b[m utf8parse v0.1.0 +\x1b[m\x1b[m\x1b[32m\x1b[1m Compiling\x1b[m vte v0.3.2 +\x1b[m\x1b[m\x1b[32m\x1b[1m Compiling\x1b[m strip-ansi-escapes v0.1.0-pre (file:///build/strip-ansi-escapes) +\x1b[m\x1b[m\x1b[32m\x1b[1m Finished\x1b[m dev [unoptimized + debuginfo] target(s) in 0.66 secs +", + b" Compiling utf8parse v0.1.0 + Compiling vte v0.3.2 + Compiling strip-ansi-escapes v0.1.0-pre (file:///build/strip-ansi-escapes) + Finished dev [unoptimized + debuginfo] target(s) in 0.66 secs +"); + } +}