1#!/usr/bin/env python3 2# Copyright 2022 The ChromiumOS Authors 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6import os 7import re 8import subprocess 9from concurrent.futures import ThreadPoolExecutor 10from dataclasses import dataclass 11from fnmatch import fnmatch 12from pathlib import Path 13import sys 14from time import sleep 15from typing import Callable, List, NamedTuple, Optional, Set, Union 16from datetime import datetime, timedelta 17 18from impl.common import Command, all_tracked_files, cmd, console, verbose 19 20import rich 21import rich.console 22import rich.live 23import rich.spinner 24import rich.text 25 26git = cmd("git") 27 28ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])") 29 30 31@dataclass 32class CheckContext(object): 33 "Information passed to each check when it's called." 34 35 # Whether or not --fix was set and checks should attempt to fix problems they encounter. 36 fix: bool 37 38 # Use rust nightly version for rust checks 39 nightly_fmt: bool 40 41 # All files that this check should cover (e.g. all python files on a python check). 42 all_files: List[Path] 43 44 # Those files of all_files that were modified locally. 45 modified_files: List[Path] 46 47 # Files that do not exist upstream and have been added locally. 48 new_files: List[Path] 49 50 51class Check(NamedTuple): 52 "Metadata for each check, definining on which files it should run." 53 54 # Function to call for this check 55 check_function: Callable[[CheckContext], Union[Command, None, List[Command]]] 56 57 custom_name: Optional[str] = None 58 59 # List of globs that this check should be triggered on 60 files: List[str] = [] 61 62 python_tools: bool = False 63 64 # List of globs to exclude from this check 65 exclude: List[str] = [] 66 67 # Whether or not this check can fix issues. 68 can_fix: bool = False 69 70 # Which groups this check belongs to. 71 groups: List[str] = [] 72 73 # Priority tasks usually take lonkger and are started first, and will show preliminary output. 74 priority: bool = False 75 76 @property 77 def name(self): 78 if self.custom_name: 79 return self.custom_name 80 name = self.check_function.__name__ 81 if name.startswith("check_"): 82 return name[len("check_") :] 83 return name 84 85 @property 86 def doc(self): 87 if self.check_function.__doc__: 88 return self.check_function.__doc__.strip() 89 else: 90 return None 91 92 93class Group(NamedTuple): 94 "Metadata for a group of checks" 95 96 name: str 97 98 doc: str 99 100 checks: List[str] 101 102 103def list_file_diff(): 104 """ 105 Lists files there were modified compared to the upstream branch. 106 107 Falls back to all files tracked by git if there is no upstream branch. 108 """ 109 upstream = git("rev-parse @{u}").stdout(check=False) 110 if upstream: 111 for line in git("diff --name-status", upstream).lines(): 112 parts = line.split("\t", 1) 113 file = Path(parts[1].strip()) 114 if file.is_file(): 115 yield (parts[0].strip(), file) 116 else: 117 print("WARNING: Not tracking a branch. Checking all files.") 118 for file in all_tracked_files(): 119 yield ("M", file) 120 121 122def should_run_check_on_file(check: Check, file: Path): 123 "Returns true if `file` should be run on `check`." 124 125 # Skip third_party except vmm_vhost. 126 if str(file).startswith("third_party") and not str(file).startswith("third_party/vmm_vhost"): 127 return False 128 129 # Skip excluded files 130 for glob in check.exclude: 131 if fnmatch(str(file), glob): 132 return False 133 134 # Match python tools (no file-extension, but with a python shebang line) 135 if check.python_tools: 136 if fnmatch(str(file), "tools/*") and file.suffix == "" and file.is_file(): 137 if file.open(errors="ignore").read(32).startswith("#!/usr/bin/env python3"): 138 return True 139 140 # If no constraint is specified, match all files. 141 if not check.files and not check.python_tools: 142 return True 143 144 # Otherwise, match only those specified by `files`. 145 for glob in check.files: 146 if fnmatch(str(file), glob): 147 return True 148 149 return False 150 151 152class Task(object): 153 """ 154 Represents a task that needs to be executed to perform a `Check`. 155 156 The task can be executed via `Task.execute`, which will update the state variables with 157 status and progress information. 158 159 This information can then be rendered from a separate thread via `Task.status_widget()` 160 """ 161 162 def __init__(self, title: str, commands: List[Command], priority: bool): 163 "Display title." 164 self.title = title 165 "Commands to execute." 166 self.commands = commands 167 "Task is a priority check." 168 self.priority = priority 169 "List of log lines (stdout+stderr) produced by the task." 170 self.log_lines: List[str] = [] 171 "Task was compleded, but may or not have been successful." 172 self.done = False 173 "True if the task completed successfully." 174 self.success = False 175 "Time the task was started." 176 self.start_time = datetime.min 177 "Duration the task took to execute. Only filled after completion." 178 self.duration = timedelta.max 179 "Spinner object for status_widget UI." 180 self.spinner = rich.spinner.Spinner("point", title) 181 182 def status_widget(self): 183 "Returns a rich console object showing the currrent status of the task." 184 duration = self.duration if self.done else datetime.now() - self.start_time 185 title = f"[{duration.total_seconds():6.2f}s] [bold]{self.title}[/bold]" 186 187 if self.done: 188 status: str = "[green]OK [/green]" if self.success else "[red]ERR[/red]" 189 title_widget = rich.text.Text.from_markup(f"{status} {title}") 190 else: 191 self.spinner.text = rich.text.Text.from_markup(title) 192 title_widget = self.spinner 193 194 if not self.priority: 195 return title_widget 196 197 last_lines = [ 198 self.log_lines[-3] if len(self.log_lines) >= 3 else "", 199 self.log_lines[-2] if len(self.log_lines) >= 2 else "", 200 self.log_lines[-1] if len(self.log_lines) >= 1 else "", 201 ] 202 203 return rich.console.Group( 204 *( 205 # Print last log lines without it's original colors 206 rich.text.Text( 207 "│ " + ansi_escape.sub("", log_line), 208 style="light_slate_grey", 209 overflow="ellipsis", 210 no_wrap=True, 211 ) 212 for log_line in last_lines 213 ), 214 rich.text.Text("└ ", end="", style="light_slate_grey"), 215 title_widget, 216 rich.text.Text(), 217 ) 218 219 def execute(self): 220 "Execute the task while updating the status variables." 221 self.start_time = datetime.now() 222 success = True 223 for command in self.commands: 224 if verbose(): 225 self.log_lines.append(f"$ {command}") 226 process = command.popen(stderr=subprocess.STDOUT) 227 assert process.stdout 228 for line in iter(process.stdout.readline, ""): 229 self.log_lines.append(line.strip()) 230 if process.wait() != 0: 231 success = False 232 self.duration = datetime.now() - self.start_time 233 self.success = success 234 self.done = True 235 236 237def print_logs(tasks: List[Task]): 238 "Prints logs of all failed or unfinished tasks." 239 for task in tasks: 240 if not task.done: 241 print() 242 console.rule(f"{task.title} did not finish", style="yellow") 243 for line in task.log_lines: 244 print(line) 245 if not task.log_lines: 246 print(f"{task.title} did not output any logs") 247 for task in tasks: 248 if task.done and not task.success: 249 console.rule(f"{task.title} failed", style="red") 250 for line in task.log_lines: 251 print(line) 252 if not task.log_lines: 253 print(f"{task.title} did not output any logs") 254 255 256def print_summary(tasks: List[Task]): 257 "Prints a summary of all task results." 258 console.rule("Summary") 259 tasks.sort(key=lambda t: t.duration) 260 for task in tasks: 261 title = f"[{task.duration.total_seconds():6.2f}s] [bold]{task.title}[/bold]" 262 status: str = "[green]OK [/green]" if task.success else "[red]ERR[/red]" 263 console.print(f"{status} {title}") 264 265 266def execute_tasks_parallel(tasks: List[Task]): 267 "Executes the list of tasks in parallel, while rendering live status updates." 268 with ThreadPoolExecutor() as executor: 269 try: 270 # Since tasks are executed in subprocesses, we can use a thread pool to parallelize 271 # despite the GIL. 272 task_futures = [executor.submit(lambda: t.execute()) for t in tasks] 273 274 # Render task updates while they are executing in the background. 275 with rich.live.Live(refresh_per_second=30) as live: 276 while True: 277 live.update( 278 rich.console.Group( 279 *(t.status_widget() for t in tasks), 280 rich.text.Text(), 281 rich.text.Text.from_markup( 282 "[green]Tip:[/green] Press CTRL-C to abort execution and see all logs." 283 ), 284 ) 285 ) 286 if all(future.done() for future in task_futures): 287 break 288 sleep(0.1) 289 except KeyboardInterrupt: 290 print_logs(tasks) 291 # Force exit to skip waiting for the executor to shutdown. This will kill all 292 # running subprocesses. 293 os._exit(1) # type: ignore 294 295 # Render error logs and summary after execution 296 print_logs(tasks) 297 print_summary(tasks) 298 299 if any(not t.success for t in tasks): 300 raise Exception("Some checks failed") 301 302 303def execute_tasks_serial(tasks: List[Task]): 304 "Executes the list of tasks one-by-one" 305 for task in tasks: 306 console.rule(task.title) 307 for command in task.commands: 308 command.fg() 309 console.print() 310 311 312def generate_plan( 313 checks_list: List[Check], 314 fix: bool, 315 run_on_all_files: bool, 316 nightly_fmt: bool, 317): 318 "Generates a list of `Task`s to execute the checks provided in `checks_list`" 319 all_files = [*all_tracked_files()] 320 file_diff = [*list_file_diff()] 321 new_files = [f for (s, f) in file_diff if s == "A"] 322 if run_on_all_files: 323 modified_files = all_files 324 else: 325 modified_files = [f for (s, f) in file_diff if s in ("M", "A")] 326 327 tasks: List[Task] = [] 328 unsupported_checks: List[str] = [] 329 for check in checks_list: 330 if fix and not check.can_fix: 331 continue 332 context = CheckContext( 333 fix=fix, 334 nightly_fmt=nightly_fmt, 335 all_files=[f for f in all_files if should_run_check_on_file(check, f)], 336 modified_files=[f for f in modified_files if should_run_check_on_file(check, f)], 337 new_files=[f for f in new_files if should_run_check_on_file(check, f)], 338 ) 339 if context.modified_files: 340 commands = check.check_function(context) 341 if commands is None: 342 unsupported_checks.append(check.name) 343 continue 344 if not isinstance(commands, list): 345 commands = [commands] 346 title = f"fixing {check.name}" if fix else check.name 347 tasks.append(Task(title, commands, check.priority)) 348 349 if unsupported_checks: 350 console.print("[yellow]Warning:[/yellow] The following checks cannot be run:") 351 for unsupported_check in unsupported_checks: 352 console.print(f" - {unsupported_check}") 353 console.print() 354 console.print("[green]Tip:[/green] Use the dev container to run presubmits:") 355 console.print() 356 console.print( 357 f" [blue] $ tools/dev_container tools/presubmit {' '.join(sys.argv[1:])}[/blue]" 358 ) 359 console.print() 360 361 # Sort so that priority tasks are launched (and rendered) first 362 tasks.sort(key=lambda t: (t.priority, t.title), reverse=True) 363 return tasks 364 365 366def run_checks( 367 checks_list: List[Check], 368 fix: bool, 369 run_on_all_files: bool, 370 nightly_fmt: bool, 371 parallel: bool, 372): 373 """ 374 Runs all checks in checks_list. 375 376 Arguments: 377 fix: Run fixes instead of checks on `Check`s that support it. 378 run_on_all_files: Do not use git delta, but run on all files. 379 nightly_fmt: Use nightly version of rust tooling. 380 parallel: Run tasks in parallel. 381 """ 382 tasks = generate_plan(checks_list, fix, run_on_all_files, nightly_fmt) 383 if len(tasks) == 1: 384 parallel = False 385 386 if parallel: 387 execute_tasks_parallel(list(tasks)) 388 else: 389 execute_tasks_serial(list(tasks)) 390