diff --git a/index.ts b/index.ts index 9569354..6d4d30e 100644 --- a/index.ts +++ b/index.ts @@ -169,6 +169,9 @@ export class PythonShell extends EventEmitter { let self = this; let errorData = ''; + // A splitter can finish before or after its source stream. + let stdoutSplitterHasEnded = true; + let stderrSplitterHasEnded = true; EventEmitter.call(this); options = extend({}, PythonShell.defaultOptions, options); @@ -201,22 +204,32 @@ export class PythonShell extends EventEmitter { // for example JSON parsing breaks if it recieves partial JSON // so we use newlineTransformer to emit each batch seperated by newline if (this.parser && this.stdout) { + stdoutSplitterHasEnded = false; if (!stdoutSplitter) stdoutSplitter = new NewlineTransformer(); // note that setting the encoding turns the chunk into a string stdoutSplitter.setEncoding(options.encoding || 'utf8'); this.stdout.pipe(stdoutSplitter).on('data', (chunk: string) => { this.emit('message', self.parser(chunk)); }); + stdoutSplitter.on('end', () => { + stdoutSplitterHasEnded = true; + terminateIfNeeded(); + }); } // listen to stderr and emit errors for incoming data if (this.stderrParser && this.stderr) { + stderrSplitterHasEnded = false; if (!stderrSplitter) stderrSplitter = new NewlineTransformer(); // note that setting the encoding turns the chunk into a string stderrSplitter.setEncoding(options.encoding || 'utf8'); this.stderr.pipe(stderrSplitter).on('data', (chunk: string) => { this.emit('stderr', self.stderrParser(chunk)); }); + stderrSplitter.on('end', () => { + stderrSplitterHasEnded = true; + terminateIfNeeded(); + }); } if (this.stderr) { @@ -253,6 +266,8 @@ export class PythonShell extends EventEmitter { if ( !self.stderrHasEnded || !self.stdoutHasEnded || + !stdoutSplitterHasEnded || + !stderrSplitterHasEnded || (self.exitCode == null && self.exitSignal == null) ) return; diff --git a/test/test-python-shell.ts b/test/test-python-shell.ts index 3d5de81..b3ed7ec 100644 --- a/test/test-python-shell.ts +++ b/test/test-python-shell.ts @@ -1,8 +1,9 @@ import * as should from 'should'; -import { PythonShell } from '..'; +import { NewlineTransformer, PythonShell } from '..'; import { sep, join } from 'path'; import { EOL as newline } from 'os'; import { chdir, cwd } from 'process'; +import { Transform } from 'stream'; describe('PythonShell', function () { const pythonFolder = 'test/python'; @@ -535,6 +536,118 @@ describe('PythonShell', function () { }); describe('.end(callback)', function () { + for (const stream of ['stdout', 'stderr']) { + for (const phase of ['transform', 'flush']) { + for (const exitCode of [0, 7]) { + it(`should wait for asynchronous ${stream} ${phase} before completing with code ${exitCode}`, async function () { + let release: () => void; + const childClosed = new Promise((resolve) => { + release = resolve; + }); + class AsyncSplitter extends NewlineTransformer { + _transform(chunk, encoding, callback) { + if (phase === 'transform') { + childClosed.then(() => + super._transform(chunk, encoding, callback), + ); + } else { + super._transform(chunk, encoding, callback); + } + } + _flush(callback) { + if (phase === 'flush') { + childClosed.then(() => super._flush(callback)); + } else { + super._flush(callback); + } + } + } + const splitter = new AsyncSplitter(); + const output = phase === 'transform' ? 'tail\n' : 'tail'; + // Hold processing until the real child has closed, without timing assumptions. + const pyshell = new PythonShell( + '-c', + { + scriptPath: '', + args: [ + `import sys; sys.${stream}.write(${JSON.stringify(output)}); sys.exit(${exitCode})`, + ], + }, + stream === 'stdout' ? splitter : null, + stream === 'stderr' ? splitter : null, + ); + pyshell.childProcess.once('close', release); + const events = []; + pyshell.on(stream === 'stdout' ? 'message' : 'stderr', (data) => { + events.push(data); + }); + pyshell.on('close', () => events.push('close')); + const splitterEnded = new Promise((resolve, reject) => { + splitter.once('end', resolve); + splitter.once('error', reject); + }); + const completed = new Promise((resolve, reject) => { + pyshell.once('error', reject); + pyshell.end((err, code) => { + try { + code.should.equal(exitCode); + if (exitCode) { + err.exitCode.should.equal(exitCode); + err.message.should.equal( + stream === 'stderr' + ? output.replace(/\n/g, newline) + : 'process exited with code 7', + ); + } else { + should.not.exist(err); + } + events.push('end'); + resolve(); + } catch (error) { + reject(error); + } + }); + }); + await Promise.all([completed, splitterEnded]); + events.should.eql(['tail', 'close', 'end']); + }); + } + } + } + + it('should collect all stderr when its splitter ends before the source', function (done) { + const splitter = new Transform({ + transform(chunk, encoding, callback) { + this.push(null); + callback(); + }, + }); + const pyshell = new PythonShell( + '-c', + { + scriptPath: '', + args: [ + 'import sys; sys.stderr.write("prefix"); sys.stderr.flush(); ' + + 'sys.stdin.readline(); sys.stderr.write("tail"); sys.exit(7)', + ], + }, + null, + splitter, + ); + pyshell.once('error', done); + pyshell.stderr.once('data', () => { + pyshell.stderr.pause(); + pyshell.send('continue').end((err) => { + err.message.should.equal('prefixtail'); + err.exitCode.should.equal(7); + done(); + }); + }); + pyshell.childProcess.once('exit', () => { + setImmediate(() => pyshell.stderr.resume()); + }); + }); + it('should end normally when exit code is zero', function (done) { let pyshell = new PythonShell('exit-code.py'); pyshell.end(function (err, code, signal) {