Skip to content

Commit 3cfc8be

Browse files
committed
Use subprocess.run instead of subprocess.Popen
(setting env variables with the env parameter)
1 parent e458cf0 commit 3cfc8be

File tree

1 file changed

+20
-27
lines changed

1 file changed

+20
-27
lines changed

src/boututils/run_wrapper.py

Lines changed: 20 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ def getmpirun(default=DEFAULT_MPIRUN):
3232
return MPIRUN
3333

3434

35-
def shell(command, pipe=False):
35+
def shell(command, pipe=False, env=None):
3636
"""Run a shell command
3737
3838
Parameters
@@ -48,24 +48,18 @@ def shell(command, pipe=False):
4848
tuple : (int, str)
4949
The return code, and either command output if pipe=True else None
5050
"""
51-
output = None
52-
status = 0
53-
if pipe:
54-
child = Popen(command, stderr=STDOUT, stdout=PIPE, shell=True)
55-
# This returns a b'string' which is casted to string in
56-
# python 2. However, as we want to use f.write() in our
57-
# runtest, we cast this to utf-8 here
58-
output = child.stdout.read().decode("utf-8", "ignore")
59-
# Wait for the process to finish. Note that child.wait()
60-
# would have deadlocked the system as stdout is PIPEd, we
61-
# therefore use communicate, which in the end also waits for
62-
# the process to finish
63-
child.communicate()
64-
status = child.returncode
65-
else:
66-
status = call(command, shell=True)
67-
68-
return status, output
51+
52+
result = subprocess.run(
53+
command,
54+
shell=True,
55+
capture_output=pipe,
56+
env=env,
57+
text=True
58+
)
59+
output = result.stdout if pipe else ''
60+
if result.stderr:
61+
output += '\nSTDERR:\n' + result.stderr
62+
return result.returncode, output
6963

7064

7165
def determineNumberOfCPUs():
@@ -236,17 +230,15 @@ def launch(
236230
if output is not None:
237231
cmd = f"{cmd} > {output}"
238232

239-
if mthread is not None:
240-
if os.name == "nt":
241-
# We're on windows, so we have to do it a little different
242-
cmd = f'cmd /C "set OMP_NUM_THREADS={mthread} && {cmd}"'
243-
else:
244-
cmd = f"OMP_NUM_THREADS={mthread} {cmd}"
233+
# Set OMP_NUM_THREADS if mthread is provided (for OpenMP in BOUT++)
234+
env = dict(os.environ)
235+
if mthread:
236+
env["OMP_NUM_THREADS"] = str(mthread)
245237

246238
if verbose:
247239
print(cmd)
248240

249-
return shell(cmd, pipe=pipe)
241+
return shell(cmd, pipe=pipe, env=env)
250242

251243

252244
def shell_safe(command, *args, **kwargs):
@@ -285,8 +277,9 @@ def launch_safe(command, *args, **kwargs):
285277
Optional arguments passed to `shell`
286278
287279
"""
280+
288281
s, out = launch(command, *args, **kwargs)
289-
if s:
282+
if s != 0:
290283
raise RuntimeError(
291284
f"Run failed with {s}.\nCommand was:\n{command}\n\nOutput was\n\n{out}"
292285
)

0 commit comments

Comments
 (0)