Skip to content

Commit 4bb470b

Browse files
authored
feat: let $ be piped from stream (google#953)
continues google#949
1 parent 8900e45 commit 4bb470b

File tree

4 files changed

+101
-37
lines changed

4 files changed

+101
-37
lines changed

.size-limit.json

+3-3
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
{
33
"name": "zx/core",
44
"path": ["build/core.cjs", "build/util.cjs", "build/vendor-core.cjs"],
5-
"limit": "72 kB",
5+
"limit": "73 kB",
66
"brotli": false,
77
"gzip": false
88
},
@@ -16,7 +16,7 @@
1616
{
1717
"name": "dts libdefs",
1818
"path": "build/*.d.ts",
19-
"limit": "36 kB",
19+
"limit": "37 kB",
2020
"brotli": false,
2121
"gzip": false
2222
},
@@ -30,7 +30,7 @@
3030
{
3131
"name": "all",
3232
"path": "build/*",
33-
"limit": "835 kB",
33+
"limit": "840 kB",
3434
"brotli": false,
3535
"gzip": false
3636
}

src/core.ts

+57-2
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,6 @@ import {
4848
once,
4949
parseDuration,
5050
preferLocalBin,
51-
promisifyStream,
5251
quote,
5352
quotePowerShell,
5453
} from './util.js'
@@ -373,7 +372,7 @@ export class ProcessPromise extends Promise<ProcessOutput> {
373372
return dest
374373
}
375374
from.once('end', () => dest.emit('end-piped-from')).pipe(dest)
376-
return promisifyStream(dest)
375+
return promisifyStream(dest, this)
377376
}
378377

379378
abort(reason?: string) {
@@ -544,6 +543,32 @@ export class ProcessPromise extends Promise<ProcessOutput> {
544543
): Promise<ProcessOutput | T> {
545544
return super.catch(onrejected)
546545
}
546+
547+
// Stream-like API
548+
private writable = true
549+
private emit(event: string, ...args: any[]) {
550+
return this
551+
}
552+
private on(event: string, cb: any) {
553+
this._stdin.on(event, cb)
554+
return this
555+
}
556+
private once(event: string, cb: any) {
557+
this._stdin.once(event, cb)
558+
return this
559+
}
560+
private write(data: any, encoding: BufferEncoding, cb: any) {
561+
this._stdin.write(data, encoding, cb)
562+
return this
563+
}
564+
private end(chunk: any, cb: any) {
565+
this._stdin.end(chunk, cb)
566+
return this
567+
}
568+
private removeListener(event: string, cb: any) {
569+
this._stdin.removeListener(event, cb)
570+
return this
571+
}
547572
}
548573

549574
type GettersRecord<T extends Record<any, any>> = { [K in keyof T]: () => T[K] }
@@ -841,3 +866,33 @@ export function log(entry: LogEntry) {
841866
process.stderr.write(entry.error + '\n')
842867
}
843868
}
869+
870+
export const promisifyStream = <S extends Writable>(
871+
stream: S,
872+
from?: ProcessPromise
873+
): S & PromiseLike<S> =>
874+
new Proxy(stream as S & PromiseLike<S>, {
875+
get(target, key) {
876+
if (key === 'run') return from?.run.bind(from)
877+
if (key === 'then') {
878+
return (res: any = noop, rej: any = noop) =>
879+
new Promise((_res, _rej) =>
880+
target
881+
.once('error', (e) => _rej(rej(e)))
882+
.once('finish', () => _res(res(target)))
883+
.once('end-piped-from', () => _res(res(target)))
884+
)
885+
}
886+
const value = Reflect.get(target, key)
887+
if (key === 'pipe' && typeof value === 'function') {
888+
return function (...args: any) {
889+
const piped = value.apply(target, args)
890+
piped._pipedFrom = from
891+
return piped instanceof ProcessPromise
892+
? piped
893+
: promisifyStream(piped, from)
894+
}
895+
}
896+
return value
897+
},
898+
})

src/util.ts

-25
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ import os from 'node:os'
1616
import path from 'node:path'
1717
import fs from 'node:fs'
1818
import { chalk } from './vendor-core.js'
19-
import type { Writable } from 'node:stream'
2019

2120
export { isStringLiteral } from './vendor-core.js'
2221

@@ -450,27 +449,3 @@ export const once = <T extends (...args: any[]) => any>(fn: T) => {
450449
return (result = fn(...args))
451450
}
452451
}
453-
454-
export const promisifyStream = <S extends Writable>(
455-
stream: S
456-
): S & PromiseLike<S> =>
457-
new Proxy(stream as S & PromiseLike<S>, {
458-
get(target, key) {
459-
if (key === 'then') {
460-
return (res: any = noop, rej: any = noop) =>
461-
new Promise((_res, _rej) =>
462-
target
463-
.once('error', (e) => _rej(rej(e)))
464-
.once('finish', () => _res(res(target)))
465-
.once('end-piped-from', () => _res(res(target)))
466-
)
467-
}
468-
const value = Reflect.get(target, key)
469-
if (key === 'pipe' && typeof value === 'function') {
470-
return function (...args: any) {
471-
return promisifyStream(value.apply(target, args))
472-
}
473-
}
474-
return value
475-
},
476-
})

test/core.test.js

+41-7
Original file line numberDiff line numberDiff line change
@@ -425,6 +425,13 @@ describe('core', () => {
425425
})
426426

427427
describe('supports chaining', () => {
428+
const getUpperCaseTransform = () =>
429+
new Transform({
430+
transform(chunk, encoding, callback) {
431+
callback(null, String(chunk).toUpperCase())
432+
},
433+
})
434+
428435
test('$ > $', async () => {
429436
const { stdout: o1 } = await $`echo "hello"`
430437
.pipe($`awk '{print $1" world"}'`)
@@ -466,13 +473,7 @@ describe('core', () => {
466473
const file = tempfile()
467474
const fileStream = fs.createWriteStream(file)
468475
const p = $`echo "hello"`
469-
.pipe(
470-
new Transform({
471-
transform(chunk, encoding, callback) {
472-
callback(null, String(chunk).toUpperCase())
473-
},
474-
})
475-
)
476+
.pipe(getUpperCaseTransform())
476477
.pipe(fileStream)
477478

478479
assert.ok(p instanceof WriteStream)
@@ -481,6 +482,39 @@ describe('core', () => {
481482
await fs.rm(file)
482483
})
483484

485+
test('$ halted > stream', async () => {
486+
const file = tempfile()
487+
const fileStream = fs.createWriteStream(file)
488+
const p = $({ halt: true })`echo "hello"`
489+
.pipe(getUpperCaseTransform())
490+
.pipe(fileStream)
491+
492+
assert.ok(p instanceof WriteStream)
493+
assert.ok(p.run() instanceof ProcessPromise)
494+
await p
495+
assert.equal((await p.run()).stdout, 'hello\n')
496+
assert.equal((await fs.readFile(file)).toString(), 'HELLO\n')
497+
await fs.rm(file)
498+
})
499+
500+
test('stream > $', async () => {
501+
const file = tempfile()
502+
await fs.writeFile(file, 'test')
503+
const { stdout } = await fs
504+
.createReadStream(file)
505+
.pipe(getUpperCaseTransform())
506+
.pipe($`cat`)
507+
508+
assert.equal(stdout, 'TEST')
509+
})
510+
511+
test('$ > stream > $', async () => {
512+
const p = $`echo "hello"`
513+
const { stdout } = await p.pipe(getUpperCaseTransform()).pipe($`cat`)
514+
515+
assert.equal(stdout, 'HELLO\n')
516+
})
517+
484518
test('$ > stdout', async () => {
485519
const p = $`echo 1`.pipe(process.stdout)
486520
assert.equal(await p, process.stdout)

0 commit comments

Comments
 (0)