feat: wrapStd

Author: pi0

examples/wrap-std.js

@@ -0,0 +1,14 @@
+#!/usr/bin/env node -r esm
+
+import { consola } from './utils'
+
+function foo () {
+  console.info('console foo') // eslint-disable-line no-console
+  process.stdout.write('called from stdout foo\n')
+  process.stderr.write('called from stderr foo\n')
+}
+
+consola.wrapStd()
+foo()
+consola.restoreStd()
+foo()

src/consola.js

@@ -1,8 +1,6 @@
 import Types from './types.js'
 import { isLogObj } from './utils/index.js'
 
-const originalConsole = Object.assign({}, console)
-
 let paused = false
 const queue = []
 
@@ -97,12 +95,60 @@ export default class Consola {
 
   wrapConsole () {
     for (const type in this._types) {
+      // Backup original value
+      if (!console['__' + type]) { // eslint-disable-line no-console
+        console['__' + type] = console[type] // eslint-disable-line no-console
+      }
+      // Override
       console[type] = this[type] // eslint-disable-line no-console
     }
   }
 
   restoreConsole () {
-    Object.assign(console, originalConsole)
+    for (const type in this._types) {
+      // Restore if backup is available
+      if (console['__' + type]) { // eslint-disable-line no-console
+        console[type] = console['__' + type] // eslint-disable-line no-console
+        delete console['__' + type] // eslint-disable-line no-console
+      }
+    }
+  }
+
+  wrapStd () {
+    this._wrapStream(this._stdout, 'log')
+    this._wrapStream(this._stderr, 'error')
+  }
+
+  _wrapStream (stream, type) {
+    if (!stream) {
+      return
+    }
+
+    // Backup original value
+    if (!stream.__write) {
+      stream.__write = stream.write
+    }
+
+    // Override
+    stream.write = (data) => {
+      this[type](String(data).trim())
+    }
+  }
+
+  restoreStd () {
+    this._restoreStream(this._stdout)
+    this._restoreStream(this._stderr)
+  }
+
+  _restoreStream (stream) {
+    if (!stream) {
+      return
+    }
+
+    if (stream.__write) {
+      stream.write = stream.__write
+      delete stream.__write
+    }
   }
 
   pause () {

src/utils/stream.js

@@ -1,10 +1,12 @@
 import { writeSync } from 'fs'
 
 export function writeStream (data, stream, mode = 'default') {
+  const write = stream.__write || stream.write
+
   switch (mode) {
     case 'async':
       return new Promise((resolve) => {
-        if (stream.write(data) === true) {
+        if (write.call(stream, data) === true) {
           resolve()
         } else {
           stream.once('drain', () => { resolve() })
@@ -13,6 +15,6 @@ export function writeStream (data, stream, mode = 'default') {
     case 'sync':
       return writeSync(stream.fd, data)
     default:
-      return stream.write(data)
+      return write.call(stream, data)
   }
 }