gae_py3/tools/generate_requirements_txt.py

#!/usr/bin/env python3
# Copyright (c) 2022 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import os
import subprocess
from pathlib import Path
from typing import List, Tuple

BASE_PATH = Path(__file__).parents[1]
REQUIREMENTS_IN = BASE_PATH.joinpath("requirements.in")
REQUIREMENTS_TXT = BASE_PATH.joinpath("requirements.txt")

MANUAL_PACKAGES = [
    # Packages which are required for GAE, but are not (yet) available in cipd
    # can be added here. Make sure those packages can be trusted since they
    # circumvent cipd's auditing process.
    #
    # ('google-cloud-logging', '3.0.0'),
]


def rm_prefix(value, prefix):
  if value.startswith(prefix):
    return value[len(prefix):]
  return value


def rm_suffix(value, suffix):
  if value.endswith(suffix):
    return value[:-len(suffix)]
  return value


class VPythonParser:
  # Adapted from https://chromium.googlesource.com/infra/infra/+/a0d5d3afba376 \
  # 30322e73035640f74669ce94ea5/go/src/infra/tools/vpython/utils/wheels.py#144

  def __init__(self, vpython_path):
    self.vpython_path = vpython_path

  def wheel_starts(self, line):
    return line.startswith('wheel: <')

  def in_wheel(self, line):
    return line.startswith('  name:') or line.startswith('  version:')

  def wheel_ends(self, line):
    return line.startswith('>')

  def clean_name(self, cipd_name):
    """Extract package name from cipd identifier.

    Examples:

      infra/python/wheels/werkzeug-py3 → werkzeug
      infra/python/wheels/markupsafe/${vpython_platform} → markupsafe
      infra/python/wheels/google-api-core-py2_py3 → google-api-core
    """
    name = rm_prefix(cipd_name, 'infra/python/wheels/')
    name = name.rsplit('/', 1)[0]  # Remove trailing `/${vpython_platform}`
    name = rm_suffix(name, '-py2_py3')
    return rm_suffix(name, '-py3')

  def parse(self) -> List[Tuple[str, str]]:
    """Read the .vpython3 file and create a list of pinned dependencies."""
    with open(self.vpython_path) as wheels_file:
      results = []
      for line in wheels_file:
        if self.wheel_starts(line):
          current_result = {}
        if self.in_wheel(line):
          parsed_line = line.split()
          key = parsed_line[0].rstrip(':')  # `name:` or `version:`
          value = parsed_line[1].strip('"')
          current_result[key] = value
        if self.wheel_ends(line):
          name = self.clean_name(current_result['name'])
          version = current_result['version'].lstrip('version:')
          results.append((name, version))

      return results


def generate_requirements_in():
  seen = set()
  vpython_deps = BASE_PATH.joinpath('.vpython3')
  deps = MANUAL_PACKAGES + VPythonParser(vpython_deps).parse()

  # Only add first appearance of each package
  cleaned_deps = list()
  for name, version in deps:
    if name in seen:
      continue
    seen.add(name)
    cleaned_deps.append((name, version))

  # Write dependencies to requirements.in file
  content = "\n".join(f'{name}=={version}' for name, version in cleaned_deps)
  with open(REQUIREMENTS_IN, 'w') as requirements_in:
    requirements_in.write(content + '\n')


def compile_requirements_txt():
  cmd = [
      'pip-compile',
      '--allow-unsafe',
      '--generate-hashes',
      '--no-header',
      '--no-annotate',
      '--strip-extras',
      '--output-file',
      REQUIREMENTS_TXT,
      REQUIREMENTS_IN,
  ]
  subprocess.run(cmd, capture_output=True, check=True)


def add_auto_generation_disclaimer():
  with open(REQUIREMENTS_TXT, 'r+') as txt:
    message = ('# This file is autogenerated. Add packages to .vpython3.\n' +
               '# \n' + '# To update, run:\n' + '# \n' +
               '#   python3 ./tools/generate_requirements_txt.py\n' + '# \n' +
               txt.read())

    # Write new content; we have already read the file, so we jump to the start
    # again
    txt.seek(0, 0)
    txt.write(message)


def cleanup():
  os.remove(REQUIREMENTS_IN)


if __name__ == '__main__':
  generate_requirements_in()
  compile_requirements_txt()
  add_auto_generation_disclaimer()
  cleanup()