5 files changed, 38 insertions, 27 deletions

diff --git a/util/ansi.js b/util/ansi.js
index 2ae5166..4e8abb0 100644
--- a/util/ansi.js
+++ b/util/ansi.js
@@ -154,7 +154,12 @@ export function disableAlternateScreen() {
 }
 
 export function measureColumns(text) {
-  // Returns the number of columns the given text takes.
+  // Returns the number of columns the given text takes. Accounts for escape
+  // codes (by not including them in the returned width).
+
+  if (text.includes(ESC)) {
+    text = text.replace(new RegExp(String.raw`${ESC}\[\??[0-9;]*.`, 'g'), '')
+  }
 
   return wcwidth(text)
 }
@@ -174,6 +179,33 @@ export function trimToColumns(text, cols) {
   return out
 }
 
+export function wrapToColumns(text, cols) {
+  // Wraps a string into separate lines. Returns an array of strings, for
+  // each line of the text.
+
+  const lines = []
+  const words = text.split(' ')
+
+  let curLine = words[0]
+  let curColumns = measureColumns(curLine)
+
+  for (const word of words.slice(1)) {
+    const wordColumns = measureColumns(word)
+    if (curColumns + wordColumns > cols) {
+      lines.push(curLine)
+      curLine = word
+      curColumns = wordColumns
+    } else {
+      curLine += ' ' + word
+      curColumns += 1 + wordColumns
+    }
+  }
+
+  lines.push(curLine)
+
+  return lines
+}
+
 export function isANSICommand(buffer, code = null) {
   return (
     buffer[0] === 0x1b && buffer[1] === 0x5b &&
diff --git a/util/index.js b/util/index.js
index db3d8a7..43a08bd 100644
--- a/util/index.js
+++ b/util/index.js
@@ -8,4 +8,3 @@ export {default as telchars} from './telchars.js'
 export {default as tuiApp} from './tui-app.js'
 export {default as unichars} from './unichars.js'
 export {default as waitForData} from './waitForData.js'
-export {default as wrap} from './wrap.js'
diff --git a/util/tui-app.js b/util/tui-app.js
index 2f09818..0dfd821 100644
--- a/util/tui-app.js
+++ b/util/tui-app.js
@@ -13,7 +13,7 @@ export default async function tuiApp(callback) {
 
     const flushable = new Flushable(process.stdout, true);
 
-    const root = new Root(screenInterface);
+    const root = new Root(screenInterface, flushable);
 
     const size = await screenInterface.getScreenSize();
     root.w = size.width;
diff --git a/util/waitForData.js b/util/waitForData.js
index f8d4a92..75f740e 100644
--- a/util/waitForData.js
+++ b/util/waitForData.js
@@ -1,9 +1,11 @@
 export default function waitForData(stream, cond = null) {
   return new Promise(resolve => {
-    stream.on('data', data => {
+    const listener = data => {
       if (cond ? cond(data) : true) {
         resolve(data)
+        stream.removeListener('data', listener)
       }
-    })
+    }
+    stream.on('data', listener)
   })
 }
diff --git a/util/wrap.js b/util/wrap.js
deleted file mode 100644
index 2c720c8..0000000
--- a/util/wrap.js
+++ /dev/null
@@ -1,22 +0,0 @@
-export default function wrap(str, width) {
-  // Wraps a string into separate lines. Returns an array of strings, for
-  // each line of the text.
-
-  const lines = []
-  const words = str.split(' ')
-
-  let curLine = words[0]
-
-  for (const word of words.slice(1)) {
-    if (curLine.length + word.length > width) {
-      lines.push(curLine)
-      curLine = word
-    } else {
-      curLine += ' ' + word
-    }
-  }
-
-  lines.push(curLine)
-
-  return lines
-}