Repository navigation
Expand file tree
/
Copy pathtask_runner.py
More file actions
356 lines (294 loc) · 12.6 KB
/
Copy pathtask_runner.py
File metadata and controls
356 lines (294 loc) · 12.6 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
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
# Copyright (c) Microsoft. All rights reserved.
"""Shared utilities for running Poe tasks across workspace packages.
These helpers centralize workspace discovery, selector matching, and execution
mode so the root task dispatcher and dependency tooling interpret package
filters the same way.
"""
import concurrent.futures
import contextlib
import glob
import os
import subprocess
import sys
import time
from collections.abc import Sequence
from fnmatch import fnmatch
from pathlib import Path
# On Windows, stdout defaults to cp1252 under non-interactive callers (e.g.
# prek / pre-commit hooks). Reconfigure to UTF-8 before importing rich so
# unicode glyphs like ``\u2713`` don't raise ``UnicodeEncodeError``.
if sys.platform == "win32":
for _stream in (sys.stdout, sys.stderr):
reconfigure = getattr(_stream, "reconfigure", None)
if callable(reconfigure):
with contextlib.suppress(OSError, ValueError):
reconfigure(encoding="utf-8")
from rich import print
from scripts._toml import tomllib
from scripts.tool_requirements import requirement_file
_TASK_TOOL_BUNDLES = {
"syntax-check": ("quality",),
"syntax-format": ("quality",),
"mypy": ("typing",),
"pyright": ("typing",),
}
_TEST_TASKS = {"integration-tests", "test", "test-integration"}
def _subprocess_env() -> dict[str, str]:
"""Return an environment that lets nested uv commands select the project env."""
env = os.environ.copy()
env.pop("VIRTUAL_ENV", None)
return env
def _has_dependency_group(project_root: Path, group_name: str) -> bool:
"""Return whether a project declares the requested dependency group."""
with (project_root / "pyproject.toml").open("rb") as file:
config = tomllib.load(file)
return group_name in (config.get("dependency-groups", {}) or {})
def discover_projects(workspace_pyproject_file: Path) -> list[Path]:
"""Discover all workspace projects from pyproject.toml."""
with workspace_pyproject_file.open("rb") as f:
data = tomllib.load(f)
projects = data["tool"]["uv"]["workspace"]["members"]
exclude = data["tool"]["uv"]["workspace"].get("exclude", [])
all_projects: list[Path] = []
for project in projects:
if "*" in project:
globbed = glob.glob(str(project), root_dir=workspace_pyproject_file.parent)
globbed_paths = [
Path(path) for path in globbed if (workspace_pyproject_file.parent / path / "pyproject.toml").is_file()
]
all_projects.extend(globbed_paths)
else:
all_projects.append(Path(project))
for project in exclude:
if "*" in project:
globbed = glob.glob(str(project), root_dir=workspace_pyproject_file.parent)
globbed_paths = [Path(p) for p in globbed]
all_projects = [p for p in all_projects if p not in globbed_paths]
else:
all_projects = [p for p in all_projects if p != Path(project)]
return all_projects
def extract_poe_tasks(file: Path) -> set[str]:
"""Extract poe task names from a pyproject.toml file."""
with file.open("rb") as f:
data = tomllib.load(f)
tasks = set(data.get("tool", {}).get("poe", {}).get("tasks", {}).keys())
# Check if there is an include too
include: str | None = data.get("tool", {}).get("poe", {}).get("include", None)
if include:
include_file = file.parent / include
if include_file.exists():
tasks = tasks.union(extract_poe_tasks(include_file))
return tasks
def build_work_items(projects: list[Path], task_names: list[str]) -> list[tuple[Path, str]]:
"""Build cross-product of (package, task) for packages that define the task."""
work_items: list[tuple[Path, str]] = []
for project in projects:
available_tasks = extract_poe_tasks(project / "pyproject.toml")
for task in task_names:
if task in available_tasks:
work_items.append((project, task))
return work_items
def normalize_project_filter(value: str) -> str:
"""Normalize a user-supplied workspace selector.
Strip presentation differences so short names, relative paths, and globs can
be compared with one matcher.
"""
normalized = value.strip().strip("/").replace("\\", "/")
return normalized or "."
def build_project_filter_candidates(project: Path | str, aliases: Sequence[str] = ()) -> set[str]:
"""Return accepted selector values for one workspace project.
We accept the workspace path, short package name, and any supplied aliases
so user-facing ``--package core`` stays stable even when underlying tools
still need paths or distribution names.
"""
normalized_path = normalize_project_filter(str(project))
candidates = {normalized_path}
if normalized_path == ".":
candidates.update({"./", "root"})
else:
# Accept bare short names like ``core`` alongside ``packages/core`` and
# ``./packages/core`` so callers do not have to care which form a
# downstream script prefers.
path = Path(normalized_path)
candidates.add(path.name)
candidates.add(f"./{normalized_path}")
for alias in aliases:
normalized_alias = normalize_project_filter(alias)
if normalized_alias and normalized_alias != ".":
candidates.add(normalized_alias)
return {candidate.lower() for candidate in candidates}
def project_filter_matches(project: Path | str, pattern: str, aliases: Sequence[str] = ()) -> bool:
"""Return whether a project matches a user-supplied selector or glob.
Matching happens against the normalized candidate set so CLI callers can use
the same selector vocabulary everywhere.
"""
normalized_pattern = normalize_project_filter(pattern).lower()
return any(
fnmatch(candidate, normalized_pattern) for candidate in build_project_filter_candidates(project, aliases)
)
def _run_task_subprocess(
project: Path,
task: str,
workspace_root: Path,
task_args: Sequence[str] = (),
) -> tuple[Path, str, int, str, str, float]:
"""Run a single poe task in a project directory via subprocess."""
start = time.monotonic()
cwd = workspace_root / project
command = _poe_task_command(cwd, task, task_args)
result = subprocess.run(
command,
cwd=cwd,
capture_output=True,
text=True,
env=_subprocess_env(),
)
elapsed = time.monotonic() - start
return (project, task, result.returncode, result.stdout, result.stderr, elapsed)
def _poe_task_command(
cwd: Path,
task: str,
task_args: Sequence[str],
) -> list[str]:
"""Build a pinned Poe command for one package-local task."""
bundles = _TASK_TOOL_BUNDLES.get(task, ())
command = ["uv", "run", "--project", str(cwd), "--locked", "--group", "dev"]
if task in _TEST_TASKS and _has_dependency_group(cwd, "test"):
command.extend(["--group", "test"])
for bundle in bundles:
command.extend(["--with-requirements", str(requirement_file(bundle))])
command.extend(["python", "-m", "poethepoet", "--executor", "simple", task, *task_args])
return command
def _run_sequential(
work_items: list[tuple[Path, str]],
workspace_root: Path,
task_args: Sequence[str] = (),
) -> None:
"""Run tasks sequentially with streaming output."""
for project, task in work_items:
print(f"Running task {task} in {project}")
cwd = workspace_root / project
result = subprocess.run(
_poe_task_command(cwd, task, task_args),
cwd=cwd,
env=_subprocess_env(),
)
if result.returncode:
sys.exit(result.returncode)
def _run_parallel(work_items: list[tuple[Path, str]], workspace_root: Path, task_args: Sequence[str] = ()) -> None:
"""Run all (package x task) combinations in parallel via subprocesses."""
max_workers = min(len(work_items), os.cpu_count() or 4)
failures: list[tuple[Path, str, str, str]] = []
completed = 0
total = len(work_items)
print(f"[cyan]Running {total} task(s) in parallel (max {max_workers} workers)...[/cyan]")
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = {
executor.submit(_run_task_subprocess, project, task, workspace_root, task_args): (project, task)
for project, task in work_items
}
for future in concurrent.futures.as_completed(futures):
project, task, returncode, stdout, stderr, elapsed = future.result()
completed += 1
progress = f"[{completed}/{total}]"
if returncode == 0:
print(f" [green]✓[/green] {progress} {task} in {project} ({elapsed:.1f}s)")
else:
print(f" [red]✗[/red] {progress} {task} in {project} ({elapsed:.1f}s)")
failures.append((project, task, stdout, stderr))
if failures:
print(f"\n[red]{len(failures)} task(s) failed:[/red]")
for project, task, stdout, stderr in failures:
print(f"\n[red]{'=' * 60}[/red]")
print(f"[red]FAILED: {task} in {project}[/red]")
if stdout.strip():
print(stdout)
if stderr.strip():
sys.stderr.write(stderr)
sys.exit(1)
print(f"\n[green]All {total} task(s) passed ✓[/green]")
def run_tasks(
work_items: list[tuple[Path, str]],
workspace_root: Path,
*,
sequential: bool = False,
task_args: Sequence[str] = (),
) -> None:
"""Run work items either in parallel or sequentially.
Single items use in-process PoeThePoet for streaming output.
Multiple items use parallel subprocesses by default.
"""
if not work_items:
print("[yellow]No matching tasks found in any package[/yellow]")
return
if sequential or len(work_items) == 1:
_run_sequential(work_items, workspace_root, task_args)
else:
_run_parallel(work_items, workspace_root, task_args)
def _run_command_subprocess(
label: str,
command: Sequence[str],
workspace_root: Path,
) -> tuple[str, int, str, str, float]:
"""Run a single labelled command in ``workspace_root`` and capture its output."""
start = time.monotonic()
result = subprocess.run(
command,
cwd=workspace_root,
capture_output=True,
text=True,
env=_subprocess_env(),
)
elapsed = time.monotonic() - start
return (label, result.returncode, result.stdout, result.stderr, elapsed)
def run_command_items(
command_items: list[tuple[str, Sequence[str]]],
workspace_root: Path,
*,
sequential: bool = False,
) -> None:
"""Run labelled commands using the same model as :func:`run_tasks`.
A single command streams its output live; multiple commands run in parallel
subprocesses with captured output and a ``✓``/``✗`` summary, mirroring the
pyright fan-out presentation.
"""
if not command_items:
print("[yellow]No commands to run[/yellow]")
return
if sequential or len(command_items) == 1:
for label, command in command_items:
print(f"[cyan]>> {label}[/cyan]")
result = subprocess.run(command, cwd=workspace_root, env=_subprocess_env())
if result.returncode:
sys.exit(result.returncode)
return
max_workers = min(len(command_items), os.cpu_count() or 4)
failures: list[tuple[str, str, str]] = []
completed = 0
total = len(command_items)
print(f"[cyan]Running {total} task(s) in parallel (max {max_workers} workers)...[/cyan]")
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = {
executor.submit(_run_command_subprocess, label, command, workspace_root): label
for label, command in command_items
}
for future in concurrent.futures.as_completed(futures):
label, returncode, stdout, stderr, elapsed = future.result()
completed += 1
progress = f"[{completed}/{total}]"
if returncode == 0:
print(f" [green]✓[/green] {progress} {label} ({elapsed:.1f}s)")
else:
print(f" [red]✗[/red] {progress} {label} ({elapsed:.1f}s)")
failures.append((label, stdout, stderr))
if failures:
print(f"\n[red]{len(failures)} task(s) failed:[/red]")
for label, stdout, stderr in failures:
print(f"\n[red]{'=' * 60}[/red]")
print(f"[red]FAILED: {label}[/red]")
if stdout.strip():
print(stdout)
if stderr.strip():
sys.stderr.write(stderr)
sys.exit(1)
print(f"\n[green]All {total} task(s) passed ✓[/green]")