tools/actions_local_runner.py - platform/external/pytorch - Git at Google

 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-

 import subprocess
 import sys
 import os
 import argparse
 import yaml
 import asyncio
 import shutil
 import re
 import fnmatch
 import shlex
 import configparser

 from typing import List, Dict, Any, Optional, Union, NamedTuple, Set

 REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


 class col:
     HEADER = "\033[95m"
     BLUE = "\033[94m"
     GREEN = "\033[92m"
     YELLOW = "\033[93m"
     RED = "\033[91m"
     RESET = "\033[0m"
     BOLD = "\033[1m"
     UNDERLINE = "\033[4m"


 def should_color() -> bool:
     return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()


 def color(the_color: str, text: str) -> str:
     if should_color():
         return col.BOLD + the_color + str(text) + col.RESET
     else:
         return text


 def cprint(the_color: str, text: str) -> None:
     if should_color():
         print(color(the_color, text))
     else:
         print(text)


 def git(args: List[str]) -> List[str]:
     p = subprocess.run(
         ["git"] + args,
         cwd=REPO_ROOT,
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
         check=True,
     )
     lines = p.stdout.decode().strip().split("\n")
     return [line.strip() for line in lines]


 def find_changed_files(ref_branch : str = "origin/master") -> List[str]:
     untracked = []

     for line in git(["status", "--porcelain"]):
         # Untracked files start with ??, so grab all of those
         if line.startswith("?? "):
             untracked.append(line.replace("?? ", ""))

     # Modified, unstaged
     modified = git(["diff", "--name-only"])

     # Modified, staged
     cached = git(["diff", "--cached", "--name-only"])

     # Committed
     merge_base = git(["merge-base", ref_branch, "HEAD"])[0]
     diff_with_origin = git(["diff", "--name-only", merge_base, "HEAD"])

     # De-duplicate
     all_files = set()
     for x in untracked + cached + modified + diff_with_origin:
         stripped = x.strip()
         if stripped != "" and os.path.exists(stripped):
             all_files.add(stripped)
     return list(all_files)


 def print_results(job_name: str, passed: bool, streams: List[str]) -> None:
     icon = color(col.GREEN, "✓") if passed else color(col.RED, "x")
     print(f"{icon} {color(col.BLUE, job_name)}")

     for stream in streams:
         stream = stream.strip()
         if stream != "":
             print(stream)


 class CommandResult(NamedTuple):
     passed: bool
     stdout: str
     stderr: str


 async def shell_cmd(
     cmd: Union[str, List[str]],
     env: Optional[Dict[str, Any]] = None,
     redirect: bool = True,
 ) -> CommandResult:
     if isinstance(cmd, list):
         cmd_str = " ".join(shlex.quote(arg) for arg in cmd)
     else:
         cmd_str = cmd

     proc = await asyncio.create_subprocess_shell(
         cmd_str,
         shell=True,
         cwd=REPO_ROOT,
         env=env,
         stdout=subprocess.PIPE if redirect else None,
         stderr=subprocess.PIPE if redirect else None,
         executable=shutil.which("bash"),
     )
     stdout, stderr = await proc.communicate()

     passed = proc.returncode == 0
     if not redirect:
         return CommandResult(passed, "", "")

     return CommandResult(passed, stdout.decode().strip(), stderr.decode().strip())


 class Check:
     name: str

     def __init__(self, files: Optional[List[str]], quiet: bool):
         self.quiet = quiet
         self.files = files

     async def run(self) -> bool:
         result = await self.run_helper()
         if result is None:
             return True

         streams = []
         if not result.passed:
             streams = [
                 result.stderr,
                 result.stdout,
             ]
         print_results(self.name, result.passed, streams)
         return result.passed

     async def run_helper(self) -> Optional[CommandResult]:
         if self.files is not None:
             relevant_files = self.filter_files(self.files)
             if len(relevant_files) == 0:
                 # No files, do nothing
                 return CommandResult(passed=True, stdout="", stderr="")

             return await self.quick(relevant_files)

         return await self.full()

     def filter_ext(self, files: List[str], extensions: Set[str]) -> List[str]:
         def passes(filename: str) -> bool:
             return os.path.splitext(filename)[1] in extensions

         return [f for f in files if passes(f)]

     def filter_files(self, files: List[str]) -> List[str]:
         return files

     async def quick(self, files: List[str]) -> CommandResult:
         raise NotImplementedError

     async def full(self) -> Optional[CommandResult]:
         raise NotImplementedError


 class Flake8(Check):
     name = "flake8"

     def filter_files(self, files: List[str]) -> List[str]:
         config = configparser.ConfigParser()
         config.read(os.path.join(REPO_ROOT, ".flake8"))

         excludes = re.split(r",\s*", config["flake8"]["exclude"].strip())
         excludes = [e.strip() for e in excludes if e.strip() != ""]

         def should_include(name: str) -> bool:
             for exclude in excludes:
                 if fnmatch.fnmatch(name, pat=exclude):
                     return False
                 if name.startswith(exclude) or f"./{name}".startswith(exclude):
                     return False
             return True

         files = self.filter_ext(files, {".py"})
         return [f for f in files if should_include(f)]

     async def quick(self, files: List[str]) -> CommandResult:
         return await shell_cmd(["flake8"] + files)

     async def full(self) -> CommandResult:
         return await shell_cmd(["flake8"])


 class Mypy(Check):
     name = "mypy (skipped typestub generation)"

     def filter_files(self, files: List[str]) -> List[str]:
         return self.filter_ext(files, {".py", ".pyi"})

     def env(self) -> Dict[str, Any]:
         env = os.environ.copy()
         if should_color():
             # Secret env variable: https://github.com/python/mypy/issues/7771
             env["MYPY_FORCE_COLOR"] = "1"
         return env

     async def quick(self, files: List[str]) -> CommandResult:
         return await shell_cmd(
             [sys.executable, "tools/linter/mypy_wrapper.py"]
             + [os.path.join(REPO_ROOT, f) for f in files],
             env=self.env(),
         )

     async def full(self) -> None:
         env = self.env()
         # hackily change the name
         self.name = "mypy"

         await shell_cmd(
             [
                 sys.executable,
                 "tools/actions_local_runner.py",
                 "--job",
                 "mypy",
                 "--file",
                 ".github/workflows/lint.yml",
                 "--step",
                 "Run autogen",
             ],
             redirect=False,
             env=env,
         )

         await shell_cmd(
             [
                 sys.executable,
                 "tools/actions_local_runner.py",
                 "--job",
                 "mypy",
                 "--file",
                 ".github/workflows/lint.yml",
                 "--step",
                 "Run mypy",
             ],
             redirect=False,
             env=env,
         )


 class ShellCheck(Check):
     name = "shellcheck: Run ShellCheck"

     def filter_files(self, files: List[str]) -> List[str]:
         return self.filter_ext(files, {".sh"})

     async def quick(self, files: List[str]) -> CommandResult:
         return await shell_cmd(
             ["tools/linter/run_shellcheck.sh"]
             + [os.path.join(REPO_ROOT, f) for f in files],
         )

     async def full(self) -> None:
         await shell_cmd(
             [
                 sys.executable,
                 "tools/actions_local_runner.py",
                 "--job",
                 "shellcheck",
                 "--file",
                 ".github/workflows/lint.yml",
                 "--step",
                 "Run ShellCheck",
             ],
             redirect=False,
         )


 class ClangTidy(Check):
     name = "clang-tidy: Run clang-tidy"
     common_options = [
         "--clang-tidy-exe",
         ".clang-tidy-bin/clang-tidy",
     ]

     def filter_files(self, files: List[str]) -> List[str]:
         return self.filter_ext(files, {".c", ".cc", ".cpp"})

     async def quick(self, files: List[str]) -> CommandResult:
         return await shell_cmd(
             [sys.executable, "-m", "tools.linter.clang_tidy", "--paths"]
             + [os.path.join(REPO_ROOT, f) for f in files]
             + self.common_options,
         )

     async def full(self) -> None:
         await shell_cmd(
             [sys.executable, "-m", "tools.linter.clang_tidy"] + self.common_options,
             redirect=False,
         )


 class YamlStep(Check):
     def __init__(self, step: Dict[str, Any], job_name: str, quiet: bool):
         super().__init__(files=None, quiet=quiet)
         self.step = step
         self.name = f'{job_name}: {self.step["name"]}'

     async def full(self) -> CommandResult:
         env = os.environ.copy()
         env["GITHUB_WORKSPACE"] = "/tmp"
         script = self.step["run"]

         if self.quiet:
             # TODO: Either lint that GHA scripts only use 'set -eux' or make this more
             # resilient
             script = script.replace("set -eux", "set -eu")
             script = re.sub(r"^time ", "", script, flags=re.MULTILINE)

         return await shell_cmd(script, env=env)


 def changed_files(ref_branch : str = "origin/master") -> Optional[List[str]]:
     changed_files: Optional[List[str]] = None
     try:
         changed_files = sorted(find_changed_files(ref_branch))
     except Exception:
         # If the git commands failed for some reason, bail out and use the whole list
         print(
             "Could not query git for changed files, falling back to testing all files instead",
             file=sys.stderr,
         )
         return None

     return changed_files


 def grab_specific_steps(
     steps_to_grab: List[str], job: Dict[str, Any]
 ) -> List[Dict[str, Any]]:
     relevant_steps = []
     for step in steps_to_grab:
         for actual_step in job["steps"]:
             if actual_step["name"].lower().strip() == step.lower().strip():
                 relevant_steps.append(actual_step)
                 break

     if len(relevant_steps) != len(steps_to_grab):
         raise RuntimeError(f"Missing steps:\n{relevant_steps}\n{steps_to_grab}")

     return relevant_steps


 def main() -> None:
     parser = argparse.ArgumentParser(
         description="Pull shell scripts out of GitHub actions and run them"
     )
     parser.add_argument("--file", help="YAML file with actions")
     parser.add_argument(
         "--changed-only",
         help="only run on changed files",
         action="store_true",
         default=False,
     )
     parser.add_argument("--job", help="job name", required=True)
     parser.add_argument(
         "--no-quiet", help="output commands", action="store_true", default=False
     )
     parser.add_argument("--step", action="append", help="steps to run (in order)")
     parser.add_argument("--ref_branch",
                         default="origin/master",
                         help="remote/branch used during comparison for --changed-only (default=origin/master")
     args = parser.parse_args()

     quiet = not args.no_quiet

     if args.file is None:
         # If there is no .yml file provided, fall back to the list of known
         # jobs. We use this for flake8 and mypy since they run different
         # locally than in CI due to 'make quicklint'
         if args.job not in ad_hoc_steps:
             raise RuntimeError(
                 f"Job {args.job} not found and no .yml file was provided"
             )

         files = None
         if args.changed_only:
             files = changed_files(args.ref_branch)

         checks = [ad_hoc_steps[args.job](files, quiet)]
     else:
         if args.step is None:
             raise RuntimeError("1+ --steps must be provided")

         action = yaml.safe_load(open(args.file, "r"))
         if "jobs" not in action:
             raise RuntimeError(f"top level key 'jobs' not found in {args.file}")
         jobs = action["jobs"]

         if args.job not in jobs:
             raise RuntimeError(f"job '{args.job}' not found in {args.file}")

         job = jobs[args.job]

         # Pull the relevant sections out of the provided .yml file and run them
         relevant_steps = grab_specific_steps(args.step, job)
         checks = [
             YamlStep(step=step, job_name=args.job, quiet=quiet)
             for step in relevant_steps
         ]

     loop = asyncio.get_event_loop()
     loop.run_until_complete(asyncio.gather(*[check.run() for check in checks]))


 # These are run differently locally in order to enable quicklint, so dispatch
 # out to special handlers instead of using lint.yml
 ad_hoc_steps = {
     "mypy": Mypy,
     "flake8-py3": Flake8,
     "shellcheck": ShellCheck,
     "clang-tidy": ClangTidy,
 }

 if __name__ == "__main__":
     try:
         main()
     except KeyboardInterrupt:
         pass
	#!/usr/bin/env python3
	# -- coding: utf-8 --

	import subprocess
	import sys
	import os
	import argparse
	import yaml
	import asyncio
	import shutil
	import re
	import fnmatch
	import shlex
	import configparser

	from typing import List, Dict, Any, Optional, Union, NamedTuple, Set

	REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


	class col:
	HEADER = "\033[95m"
	BLUE = "\033[94m"
	GREEN = "\033[92m"
	YELLOW = "\033[93m"
	RED = "\033[91m"
	RESET = "\033[0m"
	BOLD = "\033[1m"
	UNDERLINE = "\033[4m"


	def should_color() -> bool:
	return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()


	def color(the_color: str, text: str) -> str:
	if should_color():
	return col.BOLD + the_color + str(text) + col.RESET
	else:
	return text


	def cprint(the_color: str, text: str) -> None:
	if should_color():
	print(color(the_color, text))
	else:
	print(text)


	def git(args: List[str]) -> List[str]:
	p = subprocess.run(
	["git"] + args,
	cwd=REPO_ROOT,
	stdout=subprocess.PIPE,
	stderr=subprocess.PIPE,
	check=True,
	)
	lines = p.stdout.decode().strip().split("\n")
	return [line.strip() for line in lines]


	def find_changed_files(ref_branch : str = "origin/master") -> List[str]:
	untracked = []

	for line in git(["status", "--porcelain"]):
	# Untracked files start with ??, so grab all of those
	if line.startswith("?? "):
	untracked.append(line.replace("?? ", ""))

	# Modified, unstaged
	modified = git(["diff", "--name-only"])

	# Modified, staged
	cached = git(["diff", "--cached", "--name-only"])

	# Committed
	merge_base = git(["merge-base", ref_branch, "HEAD"])[0]
	diff_with_origin = git(["diff", "--name-only", merge_base, "HEAD"])

	# De-duplicate
	all_files = set()
	for x in untracked + cached + modified + diff_with_origin:
	stripped = x.strip()
	if stripped != "" and os.path.exists(stripped):
	all_files.add(stripped)
	return list(all_files)


	def print_results(job_name: str, passed: bool, streams: List[str]) -> None:
	icon = color(col.GREEN, "✓") if passed else color(col.RED, "x")
	print(f"{icon} {color(col.BLUE, job_name)}")

	for stream in streams:
	stream = stream.strip()
	if stream != "":
	print(stream)


	class CommandResult(NamedTuple):
	passed: bool
	stdout: str
	stderr: str


	async def shell_cmd(
	cmd: Union[str, List[str]],
	env: Optional[Dict[str, Any]] = None,
	redirect: bool = True,
	) -> CommandResult:
	if isinstance(cmd, list):
	cmd_str = " ".join(shlex.quote(arg) for arg in cmd)
	else:
	cmd_str = cmd

	proc = await asyncio.create_subprocess_shell(
	cmd_str,
	shell=True,
	cwd=REPO_ROOT,
	env=env,
	stdout=subprocess.PIPE if redirect else None,
	stderr=subprocess.PIPE if redirect else None,
	executable=shutil.which("bash"),
	)
	stdout, stderr = await proc.communicate()

	passed = proc.returncode == 0
	if not redirect:
	return CommandResult(passed, "", "")

	return CommandResult(passed, stdout.decode().strip(), stderr.decode().strip())


	class Check:
	name: str

	def __init__(self, files: Optional[List[str]], quiet: bool):
	self.quiet = quiet
	self.files = files

	async def run(self) -> bool:
	result = await self.run_helper()
	if result is None:
	return True

	streams = []
	if not result.passed:
	streams = [
	result.stderr,
	result.stdout,
	]
	print_results(self.name, result.passed, streams)
	return result.passed

	async def run_helper(self) -> Optional[CommandResult]:
	if self.files is not None:
	relevant_files = self.filter_files(self.files)
	if len(relevant_files) == 0:
	# No files, do nothing
	return CommandResult(passed=True, stdout="", stderr="")

	return await self.quick(relevant_files)

	return await self.full()

	def filter_ext(self, files: List[str], extensions: Set[str]) -> List[str]:
	def passes(filename: str) -> bool:
	return os.path.splitext(filename)[1] in extensions

	return [f for f in files if passes(f)]

	def filter_files(self, files: List[str]) -> List[str]:
	return files

	async def quick(self, files: List[str]) -> CommandResult:
	raise NotImplementedError

	async def full(self) -> Optional[CommandResult]:
	raise NotImplementedError


	class Flake8(Check):
	name = "flake8"

	def filter_files(self, files: List[str]) -> List[str]:
	config = configparser.ConfigParser()
	config.read(os.path.join(REPO_ROOT, ".flake8"))

	excludes = re.split(r",\s*", config["flake8"]["exclude"].strip())
	excludes = [e.strip() for e in excludes if e.strip() != ""]

	def should_include(name: str) -> bool:
	for exclude in excludes:
	if fnmatch.fnmatch(name, pat=exclude):
	return False
	if name.startswith(exclude) or f"./{name}".startswith(exclude):
	return False
	return True

	files = self.filter_ext(files, {".py"})
	return [f for f in files if should_include(f)]

	async def quick(self, files: List[str]) -> CommandResult:
	return await shell_cmd(["flake8"] + files)

	async def full(self) -> CommandResult:
	return await shell_cmd(["flake8"])


	class Mypy(Check):
	name = "mypy (skipped typestub generation)"

	def filter_files(self, files: List[str]) -> List[str]:
	return self.filter_ext(files, {".py", ".pyi"})

	def env(self) -> Dict[str, Any]:
	env = os.environ.copy()
	if should_color():
	# Secret env variable: https://github.com/python/mypy/issues/7771
	env["MYPY_FORCE_COLOR"] = "1"
	return env

	async def quick(self, files: List[str]) -> CommandResult:
	return await shell_cmd(
	[sys.executable, "tools/linter/mypy_wrapper.py"]
	+ [os.path.join(REPO_ROOT, f) for f in files],
	env=self.env(),
	)

	async def full(self) -> None:
	env = self.env()
	# hackily change the name
	self.name = "mypy"

	await shell_cmd(
	[
	sys.executable,
	"tools/actions_local_runner.py",
	"--job",
	"mypy",
	"--file",
	".github/workflows/lint.yml",
	"--step",
	"Run autogen",
	],
	redirect=False,
	env=env,
	)

	await shell_cmd(
	[
	sys.executable,
	"tools/actions_local_runner.py",
	"--job",
	"mypy",
	"--file",
	".github/workflows/lint.yml",
	"--step",
	"Run mypy",
	],
	redirect=False,
	env=env,
	)


	class ShellCheck(Check):
	name = "shellcheck: Run ShellCheck"

	def filter_files(self, files: List[str]) -> List[str]:
	return self.filter_ext(files, {".sh"})

	async def quick(self, files: List[str]) -> CommandResult:
	return await shell_cmd(
	["tools/linter/run_shellcheck.sh"]
	+ [os.path.join(REPO_ROOT, f) for f in files],
	)

	async def full(self) -> None:
	await shell_cmd(
	[
	sys.executable,
	"tools/actions_local_runner.py",
	"--job",
	"shellcheck",
	"--file",
	".github/workflows/lint.yml",
	"--step",
	"Run ShellCheck",
	],
	redirect=False,
	)


	class ClangTidy(Check):
	name = "clang-tidy: Run clang-tidy"
	common_options = [
	"--clang-tidy-exe",
	".clang-tidy-bin/clang-tidy",
	]

	def filter_files(self, files: List[str]) -> List[str]:
	return self.filter_ext(files, {".c", ".cc", ".cpp"})

	async def quick(self, files: List[str]) -> CommandResult:
	return await shell_cmd(
	[sys.executable, "-m", "tools.linter.clang_tidy", "--paths"]
	+ [os.path.join(REPO_ROOT, f) for f in files]
	+ self.common_options,
	)

	async def full(self) -> None:
	await shell_cmd(
	[sys.executable, "-m", "tools.linter.clang_tidy"] + self.common_options,
	redirect=False,
	)


	class YamlStep(Check):
	def __init__(self, step: Dict[str, Any], job_name: str, quiet: bool):
	super().__init__(files=None, quiet=quiet)
	self.step = step
	self.name = f'{job_name}: {self.step["name"]}'

	async def full(self) -> CommandResult:
	env = os.environ.copy()
	env["GITHUB_WORKSPACE"] = "/tmp"
	script = self.step["run"]

	if self.quiet:
	# TODO: Either lint that GHA scripts only use 'set -eux' or make this more
	# resilient
	script = script.replace("set -eux", "set -eu")
	script = re.sub(r"^time ", "", script, flags=re.MULTILINE)

	return await shell_cmd(script, env=env)


	def changed_files(ref_branch : str = "origin/master") -> Optional[List[str]]:
	changed_files: Optional[List[str]] = None
	try:
	changed_files = sorted(find_changed_files(ref_branch))
	except Exception:
	# If the git commands failed for some reason, bail out and use the whole list
	print(
	"Could not query git for changed files, falling back to testing all files instead",
	file=sys.stderr,
	)
	return None

	return changed_files


	def grab_specific_steps(
	steps_to_grab: List[str], job: Dict[str, Any]
	) -> List[Dict[str, Any]]:
	relevant_steps = []
	for step in steps_to_grab:
	for actual_step in job["steps"]:
	if actual_step["name"].lower().strip() == step.lower().strip():
	relevant_steps.append(actual_step)
	break

	if len(relevant_steps) != len(steps_to_grab):
	raise RuntimeError(f"Missing steps:\n{relevant_steps}\n{steps_to_grab}")

	return relevant_steps


	def main() -> None:
	parser = argparse.ArgumentParser(
	description="Pull shell scripts out of GitHub actions and run them"
	)
	parser.add_argument("--file", help="YAML file with actions")
	parser.add_argument(
	"--changed-only",
	help="only run on changed files",
	action="store_true",
	default=False,
	)
	parser.add_argument("--job", help="job name", required=True)
	parser.add_argument(
	"--no-quiet", help="output commands", action="store_true", default=False
	)
	parser.add_argument("--step", action="append", help="steps to run (in order)")
	parser.add_argument("--ref_branch",
	default="origin/master",
	help="remote/branch used during comparison for --changed-only (default=origin/master")
	args = parser.parse_args()

	quiet = not args.no_quiet

	if args.file is None:
	# If there is no .yml file provided, fall back to the list of known
	# jobs. We use this for flake8 and mypy since they run different
	# locally than in CI due to 'make quicklint'
	if args.job not in ad_hoc_steps:
	raise RuntimeError(
	f"Job {args.job} not found and no .yml file was provided"
	)

	files = None
	if args.changed_only:
	files = changed_files(args.ref_branch)

	checks = [ad_hoc_steps[args.job](files, quiet)]
	else:
	if args.step is None:
	raise RuntimeError("1+ --steps must be provided")

	action = yaml.safe_load(open(args.file, "r"))
	if "jobs" not in action:
	raise RuntimeError(f"top level key 'jobs' not found in {args.file}")
	jobs = action["jobs"]

	if args.job not in jobs:
	raise RuntimeError(f"job '{args.job}' not found in {args.file}")

	job = jobs[args.job]

	# Pull the relevant sections out of the provided .yml file and run them
	relevant_steps = grab_specific_steps(args.step, job)
	checks = [
	YamlStep(step=step, job_name=args.job, quiet=quiet)
	for step in relevant_steps
	]

	loop = asyncio.get_event_loop()
	loop.run_until_complete(asyncio.gather(*[check.run() for check in checks]))


	# These are run differently locally in order to enable quicklint, so dispatch
	# out to special handlers instead of using lint.yml
	ad_hoc_steps = {
	"mypy": Mypy,
	"flake8-py3": Flake8,
	"shellcheck": ShellCheck,
	"clang-tidy": ClangTidy,
	}

	if __name__ == "__main__":
	try:
	main()
	except KeyboardInterrupt:
	pass