text: # Copyright 2020 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.
  
  # Script lib to download all current files from GCS or upload any newly
  # generated ones.
  
  import hashlib
  import logging
  import multiprocessing
  import os
  import subprocess
  
  
  def _get_thread_count():
    """Gets a thread_count based on the multiprocessing.cpu_count()."""
    try:
      thread_count = multiprocessing.cpu_count()
      # cpu_count only gets the physical core count. There doesn't appear to be a
      # simple way of determining whether a CPU supports simultaneous
      # multithreading in Python, so assume that anything with 6 or more cores
      # supports it.
      if thread_count >= 6:
        thread_count *= 2
    except NotImplementedError:
      # Assume a quad core if we can't get the actual core count.
      thread_count = 4
    return thread_count
  
  
  def download(directory, filter, scenario, gcs_bucket):
    """Downloads the files from cloud as according to the given directory.
  
    Args:
      directory: A local disk directory.
      filter: A method to filter target files inside this directory.
      scenario: A scenario names that utilize this download method.
      gcs_bucket: A GCS bucket name to be used in the download.
    """
    # If someone removes a SHA1 file, we want to remove the associated cloud file
    # the next time file are updated.
    files_to_delete = _get_files_to_delete(directory, filter)
    for file_path in files_to_delete:
      os.remove(file_path)
  
    # Downloading the files can be very spammy, so only show the output if
    # something actually goes wrong.
    try:
      subprocess.check_output([
          'download_from_google_storage',
          '--bucket', gcs_bucket,
          '-d', directory,
          '-t', str(_get_thread_count()),
      ], stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as e:
      logging.error('Downloading %s in directory %s failed with error '
                    '%d: %s', scenario, directory, e.returncode, e.output)
  
  
  def _get_files_to_delete(directory, filter):
    """Returns a list of local files to delete.
  
    Args:
      directory: A local disk directory.
      filter: A method to filter target files inside the directory.
    """
    files_to_delete = []
    for f in os.listdir(directory):
      if not filter(f):
        continue
      sha1_path = os.path.join(directory, f + '.sha1')
      if not os.path.exists(sha1_path):
        files_to_delete.append(os.path.join(directory, f))
    return files_to_delete
  
  
  def upload(directory, filter, scenario, gcs_bucket, dry_run=False):
    """Uploads the files to cloud storage.
  
    Args:
      directory: A local disk directory.
      filter: A method to filter target files inside the directory.
      scenarios: A scenario that calls this upload.
      gcs_bucket: A GCS bucket name to be used in upload.
      dry_run: A flag to indicate whether it is just a dry_run or a real upload.
    """
    files_to_upload = _get_files_to_upload(directory, filter)
  
    if len(files_to_upload):
      if dry_run:
        logging.info('Will upload the following %s: %s',
                     scenario, ', '.join(files_to_upload))
        logging.info('Destination gcs bucket: %s', gcs_bucket)
        return
      subprocess.check_call([
          'upload_to_google_storage.py',
          '--bucket', gcs_bucket,
          '-t', str(_get_thread_count()),
      ] + files_to_upload)
  
  
  def _get_files_to_upload(directory, filter):
    """Returns a list of files to upload to cloud storage.
  
    Args:
      directory: A local disk directory.
      filter: A method to filter target files inside the directory.
    """
    files_to_upload = []
    for f in os.listdir(directory):
      # Skip any files that we don't care about.
      if not filter(f):
        continue
  
      file_path = os.path.join(directory, f)
      # upload_to_google_storage will upload a file even if it already exists
      # in the bucket. As an optimization, hash locally and only pass files to
      # the upload script if they don't have a matching .sha1 file already.
      sha_path = file_path + '.sha1'
      if os.path.isfile(sha_path):
        with open(sha_path) as sha_file:
          with open(file_path, 'rb') as source_file:
            h = hashlib.sha1()
            h.update(source_file.read())
            if sha_file.read() == h.hexdigest():
              continue
      files_to_upload.append(file_path)
    return files_to_upload
  
  
  