Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

only capture stdout not stdout+stderr when capture_output=True #2704

Merged
merged 2 commits into from
Sep 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/_nebari/provider/terraform.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,8 @@ def show(directory=None) -> dict:
run_terraform_subprocess(
command,
cwd=directory,
prefix="terraform",
strip_errors=True,
capture_output=True,
)
)
Expand Down
33 changes: 20 additions & 13 deletions src/_nebari/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,14 +46,18 @@ def change_directory(directory):
os.chdir(current_directory)


def run_subprocess_cmd(processargs, capture_output=False, **kwargs):
def run_subprocess_cmd(processargs, prefix=b"", capture_output=False, **kwargs):
"""Runs subprocess command with realtime stdout logging with optional line prefix."""
if "prefix" in kwargs:
line_prefix = f"[{kwargs['prefix']}]: ".encode("utf-8")
kwargs.pop("prefix")
if prefix:
line_prefix = f"[{prefix}]: ".encode("utf-8")
else:
line_prefix = b""

if capture_output:
stderr_stream = subprocess.PIPE
else:
stderr_stream = subprocess.STDOUT

timeout = 0
if "timeout" in kwargs:
timeout = kwargs.pop("timeout") # in seconds
Expand All @@ -64,7 +68,7 @@ def run_subprocess_cmd(processargs, capture_output=False, **kwargs):
processargs,
**kwargs,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
stderr=stderr_stream,
preexec_fn=os.setsid,
)
# Set timeout thread
Expand All @@ -80,8 +84,8 @@ def kill_process():
timeout_timer = threading.Timer(timeout, kill_process)
timeout_timer.start()

output = []
for line in iter(lambda: process.stdout.readline(), b""):
print_stream = process.stderr if capture_output else process.stdout
for line in iter(lambda: print_stream.readline(), b""):
full_line = line_prefix + line
if strip_errors:
full_line = full_line.decode("utf-8")
Expand All @@ -90,16 +94,19 @@ def kill_process():
) # Remove red ANSI escape code
full_line = full_line.encode("utf-8")

if capture_output:
output.append(full_line)
else:
sys.stdout.buffer.write(full_line)
sys.stdout.flush()
sys.stdout.buffer.write(full_line)
sys.stdout.flush()
print_stream.close()

output = []
if capture_output:
for line in iter(lambda: process.stdout.readline(), b""):
output.append(line)
process.stdout.close()

if timeout_timer is not None:
timeout_timer.cancel()

process.stdout.close()
exit_code = process.wait(
timeout=10
) # Should already have finished because we have drained stdout
Expand Down
Loading