tools/variations/split_variations_cmd.py

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

r"""A script to split Chrome variations into two sets.

Chrome runs with many experiments and variations (field trials) that are
randomly selected based on a configuration from a server. They lead to
different code paths and different Chrome behaviors. When a bug is caused by
one of the experiments or variations, it is useful to be able to bisect into
the set and pin-point which one is responsible.

Go to chrome://version/?show-variations-cmd, at the bottom a few commandline
switches define the current experiments and variations Chrome runs with.

Sample use:

python split_variations_cmd.py --file="variations_cmd.txt" --output-dir=".\out"

"variations_cmd.txt" is the command line switches data saved from
chrome://version/?show-variations-cmd. This command splits them into two sets.
If needed, the script can run on one set to further divide until a single
experiment/variation is pin-pointed as responsible.

Note that on Windows, directly passing the command line switches taken from
chrome://version/?show-variations-cmd to Chrome in "Command Prompt" won't work.
This is because Chrome in "Command Prompt" doesn't seem to handle
--force-fieldtrials="value"; it only handles --force-fieldtrials=value.
Run Chrome through "Windows PowerShell" instead.
"""

try:
  from urllib.parse import unquote, unquote_plus
except ImportError:
  # ToDo(crbug/1287214): Remove Exception case upon full migration to Python 3
  from urllib import unquote, unquote_plus

import collections
import os
import optparse
import sys
import urllib

_ENABLE_FEATURES_SWITCH_NAME = 'enable-features'
_DISABLE_FEATURES_SWITCH_NAME = 'disable-features'
_FORCE_FIELD_TRIALS_SWITCH_NAME = 'force-fieldtrials'
_FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME = 'force-fieldtrial-params'
_DISABLE_FIELD_TRIAL_CONFIG_SWITCH_NAME = 'disable-field-trial-config'


_Trial = collections.namedtuple('Trial', ['star', 'trial_name', 'group_name'])
_Param = collections.namedtuple('Param', ['key', 'value'])
_TrialParams = collections.namedtuple('TrialParams',
                                      ['trial_name', 'group_name', 'params'])
_Feature = collections.namedtuple('Feature', ['star', 'key', 'value'])

def _ParseForceFieldTrials(data):
  """Parses --force-fieldtrials switch value string.

  The switch value string is parsed to a list of _Trial objects.
  """
  data = data.rstrip('/')
  items = data.split('/')
  if len(items) % 2 != 0:
    raise ValueError('odd number of items in --force-fieldtrials value')
  trial_names = items[0::2]
  group_names = items[1::2]
  results = []
  for trial_name, group_name in zip(trial_names, group_names):
    star = False
    if trial_name.startswith('*'):
      star = True
      trial_name = trial_name[1:]
    results.append(_Trial(star, trial_name, group_name))
  return results


def _BuildForceFieldTrialsSwitchValue(trials):
  """Generates --force-fieldtrials switch value string.

  This is the opposite of _ParseForceFieldTrials().

  Args:
      trials: A list of _Trial objects from which the switch value string
          is generated.
  """
  return ''.join('%s%s/%s/' % (
    '*' if trial.star else '',
    trial.trial_name,
    trial.group_name
  ) for trial in trials)


def _ParseForceFieldTrialParams(data):
  """Parses --force-fieldtrial-params switch value string.

  The switch value string is parsed to a list of _TrialParams objects.

  Format is trial_name.group_name:param0/value0/.../paramN/valueN.
  """
  items = data.split(',')
  results = []
  for item in items:
    tokens = item.split(':')
    if len(tokens) != 2:
      raise ValueError('Wrong format, expected trial_name.group_name:'
                       'p0/v0/.../pN/vN, got %s' % item)
    trial_group = tokens[0].split('.')
    if len(trial_group) != 2:
      raise ValueError('Wrong format, expected trial_name.group_name, '
                       'got %s' % tokens[0])
    trial_name = trial_group[0]
    if len(trial_name) == 0 or trial_name[0] == '*':
      raise ValueError('Wrong field trail params format: %s' % item)
    group_name = trial_group[1]
    params = tokens[1].split('/')
    if len(params) < 2 or len(params) % 2 != 0:
      raise ValueError('Field trial params should be param/value pairs %s' %
                       tokens[1])
    pairs = [
      _Param(key=params[i], value=params[i + 1])
      for i in range(0, len(params), 2)
    ]
    results.append(_TrialParams(trial_name, group_name, pairs))
  return results


def _BuildForceFieldTrialParamsSwitchValue(trials):
  """Generates --force-fieldtrial-params switch value string.

  This is the opposite of _ParseForceFieldTrialParams().

  Args:
      trials: A list of _TrialParams objects from which the switch value
          string is generated.
  """
  return ','.join('%s.%s:%s' % (
    trial.trial_name,
    trial.group_name,
    '/'.join('%s/%s' % (param.key, param.value) for param in trial.params),
  ) for trial in trials)


def _ValidateForceFieldTrialsAndParams(trials, params):
  """Checks if all params have corresponding field trials specified.

  |trials| comes from --force-fieldtrials switch, |params| comes from
  --force-fieldtrial-params switch.
  """
  if len(params) > len(trials):
    raise ValueError("params size (%d) larger than trials size (%d)" %
                     (len(params), len(trials)))
  # Params are URL encoded, so we need to decode spaces as well. Unquote_plus
  # acts like unquote but also decodes + into spaces.
  # See https://docs.python.org/3/library/urllib.parse.html#urllib.parse.unquote_plus
  trial_groups = {trial.trial_name: trial.group_name for trial in trials}
  for param in params:
    trial_name = unquote_plus(param.trial_name)
    group_name = unquote_plus(param.group_name)
    if trial_name not in trial_groups:
      raise ValueError("Fail to find trial_name %s in trials" % trial_name)
    if group_name != trial_groups[trial_name]:
      raise ValueError("group_name mismatch for trial_name %s, %s vs %s" %
                       (trial_name, group_name, trial_groups[trial_name]))


def _SplitFieldTrials(trials, trial_params):
  """Splits (--force-fieldtrials, --force-fieldtrial-params) pair to two pairs.

  Note that any list in the output pairs could be empty, depending on the
  number of elements in the input lists.
  """
  middle = (len(trials) + 1) // 2
  params ={unquote_plus(trial.trial_name): trial for trial in trial_params}

  trials_first = trials[:middle]
  params_first = []
  for trial in trials_first:
    if trial.trial_name in params:
      params_first.append(params[trial.trial_name])

  trials_second = trials[middle:]
  params_second = []
  for trial in trials_second:
    if trial.trial_name in params:
      params_second.append(params[trial.trial_name])

  return [
    {_FORCE_FIELD_TRIALS_SWITCH_NAME: trials_first,
     _FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME: params_first},
    {_FORCE_FIELD_TRIALS_SWITCH_NAME: trials_second,
     _FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME: params_second},
  ]


def _ParseFeatureListString(data, is_disable):
  """Parses --enable/--disable-features switch value string.

  The switch value string is parsed to a list of _Feature objects.

  Args:
      data: --enable-features or --disable-features switch value string.
      is_disable: Parses --enable-features switch value string if True;
          --disable-features switch value string if False.
  """
  items = data.split(',')
  results = []
  for item in items:
    pair = item.split('<', 1)
    feature = pair[0]
    value = None
    if len(pair) > 1:
      value = pair[1]
    star = feature.startswith('*')
    if star:
      if is_disable:
        raise ValueError('--disable-features should not mark a feature with *')
      feature = feature[1:]
    results.append(_Feature(star, feature, value))
  return results


def _BuildFeaturesSwitchValue(features):
  """Generates --enable/--disable-features switch value string.

  This function does the opposite of _ParseFeatureListString().

  Args:
      features: A list of _Feature objects from which the switch value string
          is generated.
  """
  return ','.join('%s%s%s' % (
    '*' if feature.star else '',
    feature.key,
    '<%s' % feature.value if feature.value is not None else '',
  ) for feature in features)


def _SplitFeatures(features, field_trials_splits=None):
  """Splits a list of _Features objects into two lists.

  This function ensures that if a feature is associated with a specific field
  trial (by value), it stays in the same split group as that trial.

  Args:
      features: A list of _Feature objects to be split.
      field_trials_splits: A list of <=2 dictionaries, each having key
          _FORCE_FIELD_TRIALS_SWITCH_NAME mapping to a list of _Trial objects.
          This is used to keep features associated with field trials in the
          same split.

  Note that either returned list could be empty, depending on the number of
  elements in the input list, and the organized input field_trials_splits.
  """
  # Format: [ [_Feature1, _Feature2, ...], [_Feature3, ...], ... ]
  feature_groups = []
  # Format: { feature_value1:group_index1, feature_value2: group_index2, ... }
  value_to_group_index = {}

  # Group features by value, those without value gets assigned to a new group.
  for feature in features:
    if feature.value is None:
      feature_groups.append([feature])
    else:
      if feature.value in value_to_group_index:
        feature_groups[value_to_group_index[feature.value]].append(feature)
      else:
        feature_groups.append([feature])
        value_to_group_index[feature.value] = len(feature_groups) - 1

  # Extract trial names from the split field trials.
  trials_splits = []
  if field_trials_splits:
    trials_splits.append(
        set(t.trial_name for t in field_trials_splits[0].get(
            _FORCE_FIELD_TRIALS_SWITCH_NAME, [])))
    trials_splits.append(
        set(t.trial_name for t in field_trials_splits[1].get(
            _FORCE_FIELD_TRIALS_SWITCH_NAME, [])))
  else:
    trials_splits = [set(), set()]

  results = [[], []]

  # Assign features into the same split with associated field_trials.
  for group in feature_groups:
    value = group[0].value
    for i in range(2):
      if value in trials_splits[i]:
        results[i].extend(group)
        feature_groups.remove(group)
        break

  # The rest of features are split into results, balancing with best-effort.
  for group in feature_groups:
    i = 0 if results[0] <= results[1] else 1
    results[i].extend(group)

  return results[0], results[1]


def ParseCommandLineSwitchesString(data):
  """Parses command line switches string into a dictionary.

  Format: { switch1:value1, switch2:value2, ..., switchN:valueN }.
  """
  switches = data.split('--')
  # The first one is always an empty string before the first '--'.
  switches = switches[1:]
  switch_data = {}
  for switch in switches:
    switch = switch.strip()
    if switch == _DISABLE_FIELD_TRIAL_CONFIG_SWITCH_NAME:
      continue
    fields = switch.split('=', 1) # Split by the first '='.
    if len(fields) != 2:
      raise ValueError('Wrong format, expected name=value, got %s' % switch)
    switch_name, switch_value = fields
    if switch_value[0] == '"' and switch_value[-1] == '"':
      switch_value = switch_value[1:-1]
    if (switch_name == _FORCE_FIELD_TRIALS_SWITCH_NAME and
        switch_value[-1] != '/'):
      # Older versions of Chrome do not include '/' in the end, but newer
      # versions do.
      # TODO(zmo): Verify if '/' is included in the end, older versions of
      # Chrome can still accept such switch.
      switch_value = switch_value + '/'
    switch_data[switch_name] = switch_value
  return switch_data


def ParseVariationsCmdFromString(input_string):
  """Parses commandline switches string into internal representation.

  Commandline switches string comes from chrome://version/?show-variations-cmd.
  Currently we parse the following four command line switches:
    --force-fieldtrials
    --force-fieldtrial-params
    --enable-features
    --disable-features
  """
  switch_data = ParseCommandLineSwitchesString(input_string)
  results = {}
  for switch_name, switch_value in switch_data.items():
    built_switch_value = None
    if switch_name == _FORCE_FIELD_TRIALS_SWITCH_NAME:
      results[switch_name] = _ParseForceFieldTrials(switch_value)
      built_switch_value = _BuildForceFieldTrialsSwitchValue(
          results[switch_name])
    elif switch_name == _DISABLE_FEATURES_SWITCH_NAME:
      results[switch_name] = _ParseFeatureListString(
          switch_value, is_disable=True)
      built_switch_value = _BuildFeaturesSwitchValue(results[switch_name])
    elif switch_name == _ENABLE_FEATURES_SWITCH_NAME:
      results[switch_name] = _ParseFeatureListString(
          switch_value, is_disable=False)
      built_switch_value = _BuildFeaturesSwitchValue(results[switch_name])
    elif switch_name == _FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME:
      results[switch_name] = _ParseForceFieldTrialParams(switch_value)
      built_switch_value = _BuildForceFieldTrialParamsSwitchValue(
          results[switch_name])
    else:
      raise ValueError('Unexpected: --%s=%s', switch_name, switch_value)
    assert switch_value == built_switch_value
  if (_FORCE_FIELD_TRIALS_SWITCH_NAME in results and
      _FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME in results):
    _ValidateForceFieldTrialsAndParams(
        results[_FORCE_FIELD_TRIALS_SWITCH_NAME],
        results[_FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME])
  return results


def ParseVariationsCmdFromFile(filename):
  """Parses commandline switches string into internal representation.

  Same as ParseVariationsCmdFromString(), except the commandline switches
  string comes from a file.
  """
  with open(filename, 'r') as f:
    data = f.read().replace('\n', ' ')
  return ParseVariationsCmdFromString(data)


def VariationsCmdToStrings(data):
  """Converts a dictionary of {switch_name:switch_value} to a list of strings.

  Each string is in the format of '--switch_name=switch_value'.

  Args:
      data: Input data dictionary. Keys are four commandline switches:
        'force-fieldtrials'
        'force-fieldtrial-params'
        'enable-features'
        'disable-features'

  Returns:
      A list of strings.
  """
  cmd_list = []
  force_field_trials = data[_FORCE_FIELD_TRIALS_SWITCH_NAME]
  if len(force_field_trials) > 0:
    force_field_trials_switch_value = _BuildForceFieldTrialsSwitchValue(
        force_field_trials)
    cmd_list.append('--%s="%s"' % (_FORCE_FIELD_TRIALS_SWITCH_NAME,
                                   force_field_trials_switch_value))
  force_field_trial_params = data[_FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME]
  if len(force_field_trial_params) > 0:
    force_field_trial_params_switch_value = (
        _BuildForceFieldTrialParamsSwitchValue(force_field_trial_params))
    cmd_list.append('--%s="%s"' % (_FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME,
                                   force_field_trial_params_switch_value))
  enable_features = data[_ENABLE_FEATURES_SWITCH_NAME]
  if len(enable_features) > 0:
    enable_features_switch_value = _BuildFeaturesSwitchValue(enable_features)
    cmd_list.append('--%s="%s"' % (_ENABLE_FEATURES_SWITCH_NAME,
                                   enable_features_switch_value))
  disable_features = data[_DISABLE_FEATURES_SWITCH_NAME]
  if len(disable_features) > 0:
    disable_features_switch_value = _BuildFeaturesSwitchValue(
        disable_features)
    cmd_list.append('--%s="%s"' % (_DISABLE_FEATURES_SWITCH_NAME,
                                   disable_features_switch_value))
  return cmd_list


def SplitVariationsCmd(results):
  """Splits internal representation of commandline switches into two.

  This function can be called recursively when bisecting a set of experiments
  until one is identified to be responsble for a certain browser behavior.

  The commandline switches come from chrome://version/?show-variations-cmd.
  """
  enable_features = results.get(_ENABLE_FEATURES_SWITCH_NAME, [])
  disable_features = results.get(_DISABLE_FEATURES_SWITCH_NAME, [])
  field_trials = results.get(_FORCE_FIELD_TRIALS_SWITCH_NAME, [])
  field_trial_params = results.get(_FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME, [])
  field_trials_splits = _SplitFieldTrials(field_trials, field_trial_params)
  enable_features_splits = _SplitFeatures(enable_features, field_trials_splits)
  disable_features_splits = _SplitFeatures(disable_features,
                                           field_trials_splits)
  splits = []
  for index in range(2):
    cmd_line = {}
    cmd_line.update(field_trials_splits[index])
    cmd_line[_ENABLE_FEATURES_SWITCH_NAME] = enable_features_splits[index]
    cmd_line[_DISABLE_FEATURES_SWITCH_NAME] = disable_features_splits[index]
    splits.append(cmd_line)
  return splits


def SplitVariationsCmdFromString(input_string):
  """Splits commandline switches.

  This function can be called recursively when bisecting a set of experiments
  until one is identified to be responsble for a certain browser behavior.

  Same as SplitVariationsCmd(), except data comes from a string rather than
  an internal representation.

  Args:
      input_string: Variations string to be split.

  Returns:
      If input can be split, returns a list of two strings, each is half of
      the input variations cmd; otherwise, returns a list of one string.
  """
  data = ParseVariationsCmdFromString(input_string)
  splits = SplitVariationsCmd(data)
  results = []
  for split in splits:
    cmd_list = VariationsCmdToStrings(split)
    if cmd_list:
      results.append(' '.join(cmd_list))
  return results


def SplitVariationsCmdFromFile(input_filename, output_dir=None):
  """Splits commandline switches.

  This function can be called recursively when bisecting a set of experiments
  until one is identified to be responsble for a certain browser behavior.

  Same as SplitVariationsCmd(), except data comes from a file rather than
  an internal representation.

  Args:
      input_filename: Variations file to be split.
      output_dir: Folder to output the split variations file(s). If None,
          output to the same folder as the input_filename. If the folder
          doesn't exist, it will be created.

  Returns:
      If input can be split, returns a list of two output filenames;
      otherwise, returns a list of one output filename.
  """
  with open(input_filename, 'r') as f:
    input_string = f.read().replace('\n', ' ')
  splits = SplitVariationsCmdFromString(input_string)
  dirname, filename = os.path.split(input_filename)
  basename, ext = os.path.splitext(filename)
  if output_dir is None:
    output_dir = dirname
  if not os.path.exists(output_dir):
    os.makedirs(output_dir)
  split_filenames = []
  for index in range(len(splits)):
    output_filename = "%s_%d%s" % (basename, index + 1, ext)
    output_filename = os.path.join(output_dir, output_filename)
    with open(output_filename, 'w') as output_file:
      output_file.write(splits[index])
    split_filenames.append(output_filename)
  return split_filenames


def main():
  parser = optparse.OptionParser()
  parser.add_option("-f", "--file", dest="filename", metavar="FILE",
                    help="specify a file with variations cmd for processing.")
  parser.add_option("--output-dir", dest="output_dir",
                    help="specify a folder where output files are saved. "
                    "If not specified, it is the folder of the input file.")
  options, _ = parser.parse_args()
  if not options.filename:
    parser.error("Input file is not specificed")
  SplitVariationsCmdFromFile(options.filename, options.output_dir)
  return 0


if __name__ == '__main__':
  sys.exit(main())