-
-
Notifications
You must be signed in to change notification settings - Fork 106
Expand file tree
/
Copy path_shell.py
More file actions
189 lines (166 loc) · 6.07 KB
/
Copy path_shell.py
File metadata and controls
189 lines (166 loc) · 6.07 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
"""`run()`: execute a shell command per item, in parallel, with a bar.
A progress-bar'd ``xargs -P`` in Python. Templates never go through
`str.format` -- only the exact placeholder tokens ``{}`` and ``{item}``
are substituted -- so commands containing literal braces (``awk
'{print $1}'``) pass through untouched.
"""
from __future__ import annotations
import functools
import os
import shlex
import subprocess
import typing
from . import _sync
#: The two placeholder spellings recognized in command templates.
_PLACEHOLDERS: tuple[str, str] = ('{}', '{item}')
#: A command template: a string, an argv list, or a callable that
#: builds the argv for one item.
CommandT = (
str
| typing.Sequence[str]
| typing.Callable[[typing.Any], typing.Sequence[str]]
)
def _substitute(token: str, item_text: str) -> str:
"""Replace the placeholder spellings inside one token."""
for placeholder in _PLACEHOLDERS:
token = token.replace(placeholder, item_text)
return token
def _has_placeholder(tokens: typing.Iterable[str]) -> bool:
"""Return whether any token contains a placeholder."""
return any(
placeholder in token
for token in tokens
for placeholder in _PLACEHOLDERS
)
def build_argv(
command: CommandT, item: typing.Any, *, shell: bool
) -> list[str] | str:
"""Build the command for one item from a template.
Args:
command: A str template (split with `shlex.split`, non-POSIX
mode on Windows so backslash paths survive), an argv list
template, or a callable returning the argv. In the str and
list forms every ``{}``/``{item}`` inside a token is
replaced by ``str(item)`` -- an item containing spaces
stays a single argv element. Without any placeholder the
item is appended as the final argument (the ``xargs``
convention).
item: The batch item; substituted as ``str(item)``.
shell: With the str form, substitute into (and return) the
whole command string for ``subprocess.run(shell=True)``;
an appended item is `shlex.quote`-escaped. The caller must
trust its items -- see `run`.
Returns:
The argv list, or the command string when ``shell=True``.
"""
item_text: str = str(item)
if callable(command):
return [str(part) for part in command(item)]
if isinstance(command, str):
if shell:
if _has_placeholder((command,)):
return _substitute(command, item_text)
return f'{command} {shlex.quote(item_text)}'
tokens: list[str] = shlex.split(command, posix=os.name != 'nt')
else:
tokens = list(command)
if _has_placeholder(tokens):
return [_substitute(token, item_text) for token in tokens]
return [*tokens, item_text]
def _run_one(
command: CommandT,
item: typing.Any,
*,
check: bool,
capture_output: bool,
text: bool,
shell: bool,
cwd: typing.Any,
env: typing.Any,
) -> subprocess.CompletedProcess[typing.Any]:
"""Execute the command for one item (thread-pool worker)."""
argv: list[str] | str = build_argv(command, item, shell=shell)
# The argv is assembled from the caller's own template and items;
# shell=True is opt-in and documented as trusting both.
return subprocess.run( # noqa: S603, PLW1510
argv,
check=check,
capture_output=capture_output,
text=text,
shell=shell, # noqa: S602
cwd=cwd,
env=env,
)
def make_runner(
command: CommandT,
*,
check: bool = True,
capture_output: bool = True,
text: bool = True,
shell: bool = False,
cwd: typing.Any = None,
env: typing.Any = None,
) -> typing.Callable[[typing.Any], subprocess.CompletedProcess[typing.Any]]:
"""Bind a command template into a per-item callable for `map`."""
return functools.partial(
_run_one,
command,
check=check,
capture_output=capture_output,
text=text,
shell=shell,
cwd=cwd,
env=env,
)
def run(
command: CommandT,
items: typing.Iterable[typing.Any],
/,
*,
check: bool = True,
capture_output: bool = True,
text: bool = True,
shell: bool = False,
cwd: typing.Any = None,
env: typing.Any = None,
**kwargs: typing.Any,
) -> list[subprocess.CompletedProcess[typing.Any]]:
"""Run a shell command for every item in parallel, with a bar.
``progressbar.run('gzip -k {}', files, workers=4)`` is a
progress-bar'd ``xargs -P``. Subprocesses release the GIL, so this
always runs on threads (`Pool.run` reuses a pool's executor).
Args:
command: Template -- see `build_argv` for the three forms and
the placeholder rules.
items: The batch; each becomes one subprocess.
check: Raise `subprocess.CalledProcessError` on a non-zero
exit (feeding `on_error` like any other worker error).
capture_output: Capture stdout/stderr into the results --
the default, so child output cannot corrupt the bar.
text: Decode captured output as text.
shell: Run through the shell (str form only). The items are
substituted into the command line: only use with trusted
items, this is the documented injection risk.
cwd: Working directory for the subprocesses.
env: Environment for the subprocesses.
**kwargs: The shared execution keywords (`workers`, `bar`,
`on_error`, `timeout`, ...); see `_sync.execute`.
Returns:
One `subprocess.CompletedProcess` per item, in input order
(exceptions in place under ``on_error='return'``).
"""
if 'pool' in kwargs:
raise TypeError(
'run() always uses threads (subprocesses release the GIL); '
'use Pool.run() to reuse an existing pool'
)
runner = make_runner(
command,
check=check,
capture_output=capture_output,
text=text,
shell=shell,
cwd=cwd,
env=env,
)
return _sync.map(runner, items, pool='thread', **kwargs)