build/bench/utils.py

# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import importlib
import pathlib
import subprocess
import sys
import tempfile

_SRC = pathlib.Path(__file__).parent.parent.parent
_PROTO_DIR = _SRC / 'build/bench/protos'

# This ensures that we can import google.protobuf.
sys.path.append(str(_SRC / 'third_party/protobuf/python'))


def error(*args, **kwargs):
  print(*args, **kwargs, file=sys.stderr)
  exit(1)


def import_protobufs(fname: str):
  '''Imports the protobuf stored in protos/{fname}'''
  assert fname.endswith('.proto')
  path = _PROTO_DIR / fname
  with tempfile.TemporaryDirectory() as d:
    subprocess.run(
        [
            'protoc',
            f'--python_out={d}',
            f'--proto_path={_PROTO_DIR}',
            str(path),
        ],
        check=True,
    )

    sys.path.append(d)
    mod = importlib.import_module(f'{path.stem}_pb2')
    sys.path.pop()
    return mod


# We always use a capacitor file because it's self-describing, thus it can be
# directly queried without having to appear in the protodb.
class CapacitorFile:
  """A capacitor file to write to"""

  def __init__(self, out: pathlib.Path):
    # Do the validation in __init__ to prevent running a really long analysis,
    # then failing beacuse of something we could have known from the start.
    if out.suffix != '.capacitor':
      error(f'Output file must be a capacitor file. Got {out}')
    if not out.parent.resolve().is_dir():
      error(f'{out.parent} is not a directory')
    self.out = out

  def write(self, obj):
    """Writes a proto message to a capacitor file."""
    cmd = [
        'gqui',
        'from',
        'rawproto:-',
        'proto',
        f'{obj.DESCRIPTOR.file.package}.{obj.__class__.__name__}',
        f'--protofiles={_PROTO_DIR / obj.DESCRIPTOR.file.name}',
        f'--outfile=capacitor:{self.out}',
    ]
    subprocess.run(
        cmd,
        check=True,
        input=obj.SerializeToString(),
    )

  def __str__(self):
    return str(self.out)

  def __repr__(self):
    return repr(self.out)