Skip to content

Commit c3e3fb0

Browse files
committed
run_before_script: Refactor
1 parent 425bf3c commit c3e3fb0

File tree

1 file changed

+61
-30
lines changed

1 file changed

+61
-30
lines changed

src/tmuxp/util.py

Lines changed: 61 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -28,41 +28,72 @@ def run_before_script(
2828
script_file: str | pathlib.Path,
2929
cwd: pathlib.Path | None = None,
3030
) -> int:
31-
"""Execute a shell script, wraps :meth:`subprocess.check_call()` in a try/catch."""
31+
"""Execute shell script, ``tee``-ing output to both terminal (if TTY) and buffer."""
32+
script_cmd = shlex.split(str(script_file))
33+
3234
try:
3335
proc = subprocess.Popen(
34-
shlex.split(str(script_file)),
35-
stderr=subprocess.PIPE,
36-
stdout=subprocess.PIPE,
36+
script_cmd,
3737
cwd=cwd,
38-
text=True,
38+
stdout=subprocess.PIPE,
39+
stderr=subprocess.PIPE,
40+
text=True, # decode to str
3941
errors="backslashreplace",
40-
encoding="utf-8",
4142
)
42-
if proc.stdout is not None:
43-
for line in iter(proc.stdout.readline, ""):
44-
sys.stdout.write(line)
45-
proc.wait()
46-
47-
if proc.returncode and proc.stderr is not None:
48-
stderr = proc.stderr.read()
49-
proc.stderr.close()
50-
stderr_strlist = stderr.split("\n")
51-
stderr_str = "\n".join(list(filter(None, stderr_strlist))) # filter empty
52-
53-
raise exc.BeforeLoadScriptError(
54-
proc.returncode,
55-
os.path.abspath(script_file), # NOQA: PTH100
56-
stderr_str,
57-
)
58-
except OSError as e:
59-
if e.errno == 2:
60-
raise exc.BeforeLoadScriptNotExists(
61-
e,
62-
os.path.abspath(script_file), # NOQA: PTH100
63-
) from e
64-
raise
65-
return proc.returncode
43+
except FileNotFoundError as e:
44+
raise exc.BeforeLoadScriptNotExists(
45+
e,
46+
os.path.abspath(script_file), # NOQA: PTH100
47+
) from e
48+
49+
out_buffer = []
50+
err_buffer = []
51+
52+
# While process is running, read lines from stdout/stderr
53+
# and write them to this process's stdout/stderr if isatty
54+
is_out_tty = sys.stdout.isatty()
55+
is_err_tty = sys.stderr.isatty()
56+
57+
# You can do a simple loop reading in real-time:
58+
while True:
59+
# Use .poll() to check if the child has exited
60+
return_code = proc.poll()
61+
62+
# Read one line from stdout, if available
63+
line_out = proc.stdout.readline() if proc.stdout else ""
64+
65+
# Read one line from stderr, if available
66+
line_err = proc.stderr.readline() if proc.stderr else ""
67+
68+
if line_out:
69+
out_buffer.append(line_out)
70+
if is_out_tty:
71+
sys.stdout.write(line_out)
72+
sys.stdout.flush()
73+
74+
if line_err:
75+
err_buffer.append(line_err)
76+
if is_err_tty:
77+
sys.stderr.write(line_err)
78+
sys.stderr.flush()
79+
80+
# If no more data from pipes and process ended, break
81+
if not line_out and not line_err and return_code is not None:
82+
break
83+
84+
# At this point, the process has finished
85+
return_code = proc.wait()
86+
87+
if return_code != 0:
88+
# Join captured stderr lines for your exception
89+
stderr_str = "".join(err_buffer).strip()
90+
raise exc.BeforeLoadScriptError(
91+
return_code,
92+
os.path.abspath(script_file), # NOQA: PTH100
93+
stderr_str,
94+
)
95+
96+
return return_code
6697

6798

6899
def oh_my_zsh_auto_title() -> None:

0 commit comments

Comments
 (0)