From f5fe0496f52d2ee0fa1132e27555db7d8eb076a4 Mon Sep 17 00:00:00 2001 From: JustAnotherArchivist Date: Mon, 20 Mar 2023 01:13:56 +0000 Subject: [PATCH] Add support for supplying a file-like object as stdin --- codearchiver/subprocess.py | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/codearchiver/subprocess.py b/codearchiver/subprocess.py index 5a187cd..adf1b52 100644 --- a/codearchiver/subprocess.py +++ b/codearchiver/subprocess.py @@ -12,7 +12,7 @@ def run_with_log(args, *, check = True, input = None, **kwargs): ''' Run a command using `subprocess.Popen(args, **kwargs)` and log all its stderr output via `logging`. `check` has the same semantics as on `subprocess.run`, i.e. raises an exception if the process exits non-zero. - `input`, if specified, is a `bytes` object that is fed to the subprocess via stdin. + `input`, if specified, is a `bytes` or a binary file-like object that is fed to the subprocess via stdin. `stdin`, `stdout`, and `stderr` kwargs must not be used. Returns a tuple with the process's exit status, its stdout output, and its stderr output. ''' @@ -32,19 +32,30 @@ def run_with_log(args, *, check = True, input = None, **kwargs): stderr = [] stderrBuf = b'' if input is not None: - stdinView = memoryview(input) - stdinOffset = 0 + inputIsBytes = isinstance(input, bytes) + if inputIsBytes: + stdinView = memoryview(input) + stdinOffset = 0 + stdinLength = len(input) PIPE_BUF = getattr(select, 'PIPE_BUF', 512) while sel.get_map(): for key, _ in sel.select(): if key.fileobj is p.stdin: try: - stdinOffset += os.write(key.fd, stdinView[stdinOffset : stdinOffset + PIPE_BUF]) + if inputIsBytes: + stdinOffset += os.write(key.fd, stdinView[stdinOffset : stdinOffset + PIPE_BUF]) + else: + d = input.read(PIPE_BUF) + if not d: + sel.unregister(key.fileobj) + key.fileobj.close() + else: + os.write(key.fd, d) except BrokenPipeError: sel.unregister(key.fileobj) key.fileobj.close() else: - if stdinOffset >= len(input): + if inputIsBytes and stdinOffset >= stdinLength: sel.unregister(key.fileobj) key.fileobj.close() else: @@ -68,7 +79,7 @@ def run_with_log(args, *, check = True, input = None, **kwargs): _logger.info(stderrBuf.decode('utf-8')) p.wait() assert p.poll() is not None - if input is not None and stdinOffset < len(input): + if input is not None and inputIsBytes and stdinOffset < len(input): _logger.warning(f'Could not write all input to the stdin pipe (wanted to write {len(input)} bytes, only wrote {stdinOffset})') _logger.info(f'Process exited with status {p.returncode}') if check and p.returncode != 0: