scripts/run_local_regen.py

#!/usr/bin/env python3
# Copyright 2021 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Regenerate all project configs locally and generate a diff for them."""

import argparse
import atexit
import collections
import functools
import glob
import itertools
import json
import logging
import multiprocessing
import multiprocessing.pool
import os
import pathlib
import shutil
import subprocess
import sys
import tempfile
import time

from common import utilities

# resolve relative directories
this_dir = pathlib.Path(os.path.dirname(os.path.abspath(__file__)))
path_projects = this_dir / "../../project"


def run_config(config, logname=None):
  """Run a single config job, return diff of current and new output.

  Args:
    config: (program, project) to configure
    logname: Filename to redirect stderr to from config
      default is to suppress the output
  """

  # open logfile as /dev/null by default so we don't have to check for it
  logfile = open("/dev/null", "w")
  if logname:
    logfile = open(logname, "a")

  def run_diff(cmd, current, output):
    """Execute cmd and diff the current and output files"""
    stdout = ""
    try:
      stdout = subprocess.run(
          cmd,
          stdout=subprocess.PIPE,
          stderr=subprocess.STDOUT,
          check=False,
          text=True,
          shell=True,
      ).stdout
    finally:
      logfile.write(stdout)

    # otherwise run diff
    return utilities.jqdiff(
        current if current.exists() else None,
        output if output.exists() else None,
    )

  #### start of function body
  program, project = config

  logfile.write("== {}/{} ==================\n".format(program, project))

  # path to project repo and config bundle
  path_repo = (path_projects / program / project).resolve()
  path_config = (path_repo / "generated/config.jsonproto").resolve()

  # create temporary directory for output
  diff = ""
  with tempfile.TemporaryDirectory() as scratch:
    scratch = pathlib.Path(scratch)

    cmd = "cd {}; {} --no-proto --output-dir {} {}".format(
        path_repo.resolve(),
        (path_repo / "config/bin/gen_config").resolve(),
        scratch,
        path_repo / "config.star",
    )

    path_config_new = scratch / "generated/config.jsonproto"

    logfile.write("repo path:  {}\n".format(path_repo))
    logfile.write("old config: {}\n".format(path_config))
    logfile.write("new config: {}\n".format(path_config_new))
    logfile.write("running:    {}\n".format(cmd))
    logfile.write("\n")

    diff = run_diff(cmd, path_config, path_config_new)
    logfile.write("\n\n")
    logfile.close()

  return ("{}-{}".format(program, project), diff)


def run_configs(args, configs):
  """Regenerate configuration for each project in configs.

  Generate an über diff showing the changes that the current ToT
  configuration code would generate vs what's currently committed.

  Write the result to the output file specified on the command line.

  Args:
    args: command line arguments from argparse
    configs: list of BackfillConfig instances to execute

  Return:
    nothing
  """

  # create a logfile if requested
  kwargs = {}
  if args.logfile:
    # open and close the logfile to truncate it so backfills can append
    # We can't pickle the file object and send it as an argument with
    # multiprocessing, so this is a workaround for that limitation
    with open(args.logfile, "w"):
      kwargs["logname"] = args.logfile

  nproc = 32
  diffs = {}
  nconfig = len(configs)
  with multiprocessing.Pool(processes=nproc) as pool:
    results = pool.imap_unordered(
        functools.partial(run_config, **kwargs), configs, chunksize=1)
    for ii, result in enumerate(results, 1):
      sys.stderr.write(
          utilities.clear_line("[{}/{}] Processing configs".format(ii,
                                                                   nconfig)))

      if result:
        key, diff = result
        diffs[key] = diff

    sys.stderr.write(utilities.clear_line("[✔] Processing configs"))

  # generate final über diff showing all the changes
  with open(args.diff, "w") as ofile:
    for name, result in sorted(diffs.items()):
      ofile.write("## ---------------------\n")
      ofile.write("## diff for {}\n".format(name))
      ofile.write("\n")
      ofile.write(result + "\n")


def main():
  parser = argparse.ArgumentParser(
      description=__doc__,
      formatter_class=argparse.RawTextHelpFormatter,
  )

  parser.add_argument(
      "--diff",
      type=str,
      required=True,
      help="target file for diff on config.jsonproto payload",
  )

  parser.add_argument(
      "-l",
      "--logfile",
      type=str,
      help="target file to log output from regens",
  )
  args = parser.parse_args()

  # glob out all the config.stars in the project directory
  strpath = str(path_projects)
  config_stars = [
      os.path.relpath(fname, strpath)
      for fname in glob.glob(strpath + "/*/*/config.star")
  ]

  # break program/project out of path
  configs = [
      tuple(config_star.split("/")[0:2],) for config_star in config_stars
  ]
  run_configs(args, sorted(configs))


if __name__ == "__main__":
  main()