From 7eb549875db5626c362405c2ce61e39772ea9ad8 Mon Sep 17 00:00:00 2001
From: Ian Henry <ianthehenry@gmail.com>
Date: Sun, 4 Apr 2021 21:42:41 -0700
Subject: [PATCH] capture pane with escape codes

---
 Cargo.lock     | 25 +++++++++++++++++++++++++
 Cargo.toml     |  1 +
 src/main.rs    |  9 ++++++++-
 src/state.rs   |  8 +++++---
 src/swapper.rs |  2 +-
 src/view.rs    |  2 +-
 6 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 2093621..cc7d72e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -122,6 +122,15 @@ version = "0.6.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7fe5bd57d1d7414c6b5ed48563a2c855d995ff777729dcd91c369ec7fea395ae"
 
+[[package]]
+name = "strip-ansi-escapes"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9d63676e2abafa709460982ddc02a3bb586b6d15a49b75c212e06edd3933acee"
+dependencies = [
+ "vte",
+]
+
 [[package]]
 name = "strsim"
 version = "0.8.0"
@@ -166,6 +175,7 @@ dependencies = [
  "clap",
  "lazy_static",
  "regex",
+ "strip-ansi-escapes",
  "termion",
  "unicode-width",
 ]
@@ -176,12 +186,27 @@ version = "0.1.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "caaa9d531767d1ff2150b9332433f32a24622147e5ebb1f26409d5da67afd479"
 
+[[package]]
+name = "utf8parse"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8772a4ccbb4e89959023bc5b7cb8623a795caa7092d99f3aa9501b9484d4557d"
+
 [[package]]
 name = "vec_map"
 version = "0.8.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "05c78687fb1a80548ae3250346c3db86a80a7cdd77bda190189f2d0a0987c81a"
 
+[[package]]
+name = "vte"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4f42f536e22f7fcbb407639765c8fd78707a33109301f834a594758bedd6e8cf"
+dependencies = [
+ "utf8parse",
+]
+
 [[package]]
 name = "winapi"
 version = "0.3.8"
diff --git a/Cargo.toml b/Cargo.toml
index e0100f7..226a47b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,7 @@ clap = "2.33.0"
 base64 = "0.11.0"
 unicode-width = "0.1.7"
 lazy_static = "1.4.0"
+strip-ansi-escapes = "0.1.0"
 
 [[bin]]
 name = "thumbs"
diff --git a/src/main.rs b/src/main.rs
index a810286..e2374c6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -158,7 +158,14 @@ fn main() {
 
   let lines = output.split('\n').collect::<Vec<&str>>();
 
-  let mut state = state::State::new(&lines, alphabet, &regexp);
+  let unescaped = lines.iter().map(|line|
+    match strip_ansi_escapes::strip(line) {
+      Ok(bytes) => unsafe { String::from_utf8_unchecked(bytes) },
+      Err(_) => String::from("unparseable")
+    }
+  ).collect();
+
+  let mut state = state::State::new(&lines, &unescaped, alphabet, &regexp);
 
   let selected = {
     let mut viewbox = view::View::new(
diff --git a/src/state.rs b/src/state.rs
index 022c61f..02986fd 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -52,14 +52,16 @@ impl<'a> PartialEq for Match<'a> {
 
 pub struct State<'a> {
   pub lines: &'a Vec<&'a str>,
+  pub unescaped: &'a Vec<String>,
   alphabet: &'a str,
   regexp: &'a Vec<&'a str>,
 }
 
 impl<'a> State<'a> {
-  pub fn new(lines: &'a Vec<&'a str>, alphabet: &'a str, regexp: &'a Vec<&'a str>) -> State<'a> {
+  pub fn new(lines: &'a Vec<&'a str>, unescaped: &'a Vec<String>, alphabet: &'a str, regexp: &'a Vec<&'a str>) -> State<'a> {
     State {
       lines,
+      unescaped,
       alphabet,
       regexp,
     }
@@ -86,8 +88,8 @@ impl<'a> State<'a> {
 
     let all_patterns = [exclude_patterns, custom_patterns, patterns].concat();
 
-    for (index, line) in self.lines.iter().enumerate() {
-      let mut chunk: &str = line;
+    for (index, line) in self.unescaped.iter().enumerate() {
+      let mut chunk: &str = &line;
       let mut offset: i32 = 0;
 
       loop {
diff --git a/src/swapper.rs b/src/swapper.rs
index aff33e1..707d0eb 100644
--- a/src/swapper.rs
+++ b/src/swapper.rs
@@ -210,7 +210,7 @@ impl<'a> Swapper<'a> {
     };
 
     let pane_command = format!(
-        "tmux capture-pane -t {active_pane_id} -p{scroll_params} | tail -n {height} | {dir}/target/release/thumbs -f '%U:%H' -t {tmp} {args}; tmux swap-pane -t {active_pane_id}; {zoom_command} tmux wait-for -S {signal}",
+        "tmux capture-pane -et {active_pane_id} -p{scroll_params} | tail -n {height} | {dir}/target/release/thumbs -f '%U:%H' -t {tmp} {args}; tmux swap-pane -t {active_pane_id}; {zoom_command} tmux wait-for -S {signal}",
         active_pane_id = active_pane_id,
         scroll_params = scroll_params,
         height = self.active_pane_height.unwrap_or(i32::MAX),
diff --git a/src/view.rs b/src/view.rs
index ce107b6..653c5b1 100644
--- a/src/view.rs
+++ b/src/view.rs
@@ -110,7 +110,7 @@ impl<'a> View<'a> {
       };
 
       // Find long utf sequences and extract it from mat.x
-      let line = &self.state.lines[mat.y as usize];
+      let line = &self.state.unescaped[mat.y as usize];
       let prefix = &line[0..mat.x as usize];
       let offset = prefix.width() as u16;
       let text = self.make_hint_text(mat.text);