From 9124553cc406eb0544e13024437708ab57aac4f5 Mon Sep 17 00:00:00 2001
From: Anton Golub <antongolub@antongolub.com>
Date: Sun, 24 Nov 2024 18:55:18 +0300
Subject: [PATCH] fix: mixin `ProcessOutput` data to promisified pipe value

continues #949
---
 src/core.ts           | 57 ++++++++++++++++++++++---------------------
 src/util.ts           | 13 ++++++++++
 test-d/core.test-d.ts |  6 ++---
 test/core.test.js     | 15 +++++++-----
 4 files changed, 54 insertions(+), 37 deletions(-)

diff --git a/src/core.ts b/src/core.ts
index 52693aed97..9b7906fcbd 100644
--- a/src/core.ts
+++ b/src/core.ts
@@ -48,6 +48,7 @@ import {
   once,
   parseDuration,
   preferLocalBin,
+  proxyOverride,
   quote,
   quotePowerShell,
 } from './util.js'
@@ -328,12 +329,12 @@ export class ProcessPromise extends Promise<ProcessOutput> {
 
   // Essentials
   pipe(dest: TemplateStringsArray, ...args: any[]): ProcessPromise
-  pipe<D extends Writable>(dest: D): D & PromiseLike<D>
+  pipe<D extends Writable>(dest: D): D & PromiseLike<ProcessOutput & D>
   pipe<D extends ProcessPromise>(dest: D): D
   pipe(
     dest: Writable | ProcessPromise | TemplateStringsArray,
     ...args: any[]
-  ): (Writable & PromiseLike<Writable>) | ProcessPromise {
+  ): (Writable & PromiseLike<ProcessPromise & Writable>) | ProcessPromise {
     if (isStringLiteral(dest, ...args))
       return this.pipe($({ halt: true })(dest as TemplateStringsArray, ...args))
     if (isString(dest))
@@ -372,7 +373,8 @@ export class ProcessPromise extends Promise<ProcessOutput> {
       return dest
     }
     from.once('end', () => dest.emit('end-piped-from')).pipe(dest)
-    return promisifyStream(dest, this)
+    return promisifyStream(dest, this) as Writable &
+      PromiseLike<ProcessPromise & Writable>
   }
 
   abort(reason?: string) {
@@ -867,32 +869,31 @@ export function log(entry: LogEntry) {
   }
 }
 
-export const promisifyStream = <S extends Writable>(
+const promisifyStream = <S extends Writable>(
   stream: S,
-  from?: ProcessPromise
-): S & PromiseLike<S> =>
-  new Proxy(stream as S & PromiseLike<S>, {
-    get(target, key) {
-      if (key === 'run') return from?.run.bind(from)
-      if (key === 'then') {
-        return (res: any = noop, rej: any = noop) =>
-          new Promise((_res, _rej) =>
-            target
-              .once('error', (e) => _rej(rej(e)))
-              .once('finish', () => _res(res(target)))
-              .once('end-piped-from', () => _res(res(target)))
+  from: ProcessPromise
+): S & PromiseLike<ProcessOutput & S> =>
+  proxyOverride(stream as S & PromiseLike<ProcessOutput & S>, {
+    then(res: any = noop, rej: any = noop) {
+      return new Promise((_res, _rej) =>
+        stream
+          .once('error', (e) => _rej(rej(e)))
+          .once('finish', () =>
+            _res(res(proxyOverride(stream, (from as any)._output)))
           )
-      }
-      const value = Reflect.get(target, key)
-      if (key === 'pipe' && typeof value === 'function') {
-        return function (...args: any) {
-          const piped = value.apply(target, args)
-          piped._pipedFrom = from
-          return piped instanceof ProcessPromise
-            ? piped
-            : promisifyStream(piped, from)
-        }
-      }
-      return value
+          .once('end-piped-from', () =>
+            _res(res(proxyOverride(stream, (from as any)._output)))
+          )
+      )
+    },
+    run() {
+      return from.run()
+    },
+    _pipedFrom: from,
+    pipe(...args: any) {
+      const piped = stream.pipe.apply(stream, args)
+      return piped instanceof ProcessPromise
+        ? piped
+        : promisifyStream(piped as Writable, from)
     },
   })
diff --git a/src/util.ts b/src/util.ts
index 164480463d..56b7ed4890 100644
--- a/src/util.ts
+++ b/src/util.ts
@@ -449,3 +449,16 @@ export const once = <T extends (...args: any[]) => any>(fn: T) => {
     return (result = fn(...args))
   }
 }
+
+export const proxyOverride = <T extends object>(
+  origin: T,
+  ...fallbacks: any
+): T =>
+  new Proxy(origin, {
+    get(target: T, key) {
+      return (
+        fallbacks.find((f: any) => key in f)?.[key] ??
+        Reflect.get(target as T, key)
+      )
+    },
+  }) as T
diff --git a/test-d/core.test-d.ts b/test-d/core.test-d.ts
index 02bf3167aa..53567dd1f5 100644
--- a/test-d/core.test-d.ts
+++ b/test-d/core.test-d.ts
@@ -27,9 +27,9 @@ expectType<ProcessPromise>(p.nothrow())
 expectType<ProcessPromise>(p.quiet())
 expectType<ProcessPromise>(p.pipe($`cmd`))
 expectType<ProcessPromise>(p.pipe`cmd`)
-expectType<typeof process.stdout & PromiseLike<typeof process.stdout>>(
-  p.pipe(process.stdout)
-)
+expectType<
+  typeof process.stdout & PromiseLike<ProcessOutput & typeof process.stdout>
+>(p.pipe(process.stdout))
 expectType<ProcessPromise>(p.stdio('pipe'))
 expectType<ProcessPromise>(p.timeout('1s'))
 expectType<Promise<void>>(p.kill())
diff --git a/test/core.test.js b/test/core.test.js
index 715bcc9978..e6e918d4c0 100644
--- a/test/core.test.js
+++ b/test/core.test.js
@@ -475,13 +475,21 @@ describe('core', () => {
           const p = $`echo "hello"`
             .pipe(getUpperCaseTransform())
             .pipe(fileStream)
+          const o = await p
 
           assert.ok(p instanceof WriteStream)
-          assert.equal(await p, fileStream)
+          assert.ok(o instanceof WriteStream)
+          assert.equal(o.stdout, 'hello\n')
+          assert.equal(o.exitCode, 0)
           assert.equal((await fs.readFile(file)).toString(), 'HELLO\n')
           await fs.rm(file)
         })
 
+        test('$ > stdout', async () => {
+          const p = $`echo 1`.pipe(process.stdout)
+          assert.deepEqual(p, process.stdout)
+        })
+
         test('$ halted > stream', async () => {
           const file = tempfile()
           const fileStream = fs.createWriteStream(file)
@@ -514,11 +522,6 @@ describe('core', () => {
 
           assert.equal(stdout, 'HELLO\n')
         })
-
-        test('$ > stdout', async () => {
-          const p = $`echo 1`.pipe(process.stdout)
-          assert.equal(await p, process.stdout)
-        })
       })
 
       it('supports delayed piping', async () => {