michael@0: # Copyright (c) 2012 The Chromium Authors. All rights reserved.
michael@0: # Use of this source code is governed by a BSD-style license that can be
michael@0: # found in the LICENSE file.
michael@0: 
michael@0: import constants
michael@0: import traceback
michael@0: import warnings
michael@0: 
michael@0: 
michael@0: # Location where chrome reads command line flags from
michael@0: CHROME_COMMAND_FILE = constants.TEST_EXECUTABLE_DIR + '/chrome-command-line'
michael@0: 
michael@0: class FlagChanger(object):
michael@0:   """Changes the flags Chrome runs with.
michael@0: 
michael@0:   There are two different use cases for this file:
michael@0:   * Flags are permanently set by calling Set().
michael@0:   * Flags can be temporarily set for a particular set of unit tests.  These
michael@0:     tests should call Restore() to revert the flags to their original state
michael@0:     once the tests have completed.
michael@0:   """
michael@0: 
michael@0:   def __init__(self, android_cmd):
michael@0:     self._android_cmd = android_cmd
michael@0: 
michael@0:     # Save the original flags.
michael@0:     self._orig_line = self._android_cmd.GetFileContents(CHROME_COMMAND_FILE)
michael@0:     if self._orig_line:
michael@0:       self._orig_line = self._orig_line[0].strip()
michael@0: 
michael@0:     # Parse out the flags into a list to facilitate adding and removing flags.
michael@0:     self._current_flags = self._TokenizeFlags(self._orig_line)
michael@0: 
michael@0:   def Get(self):
michael@0:     """Returns list of current flags."""
michael@0:     return self._current_flags
michael@0: 
michael@0:   def Set(self, flags):
michael@0:     """Replaces all flags on the current command line with the flags given.
michael@0: 
michael@0:     Args:
michael@0:       flags: A list of flags to set, eg. ['--single-process'].
michael@0:     """
michael@0:     if flags:
michael@0:       assert flags[0] != 'chrome'
michael@0: 
michael@0:     self._current_flags = flags
michael@0:     self._UpdateCommandLineFile()
michael@0: 
michael@0:   def AddFlags(self, flags):
michael@0:     """Appends flags to the command line if they aren't already there.
michael@0: 
michael@0:     Args:
michael@0:       flags: A list of flags to add on, eg. ['--single-process'].
michael@0:     """
michael@0:     if flags:
michael@0:       assert flags[0] != 'chrome'
michael@0: 
michael@0:     # Avoid appending flags that are already present.
michael@0:     for flag in flags:
michael@0:       if flag not in self._current_flags:
michael@0:         self._current_flags.append(flag)
michael@0:     self._UpdateCommandLineFile()
michael@0: 
michael@0:   def RemoveFlags(self, flags):
michael@0:     """Removes flags from the command line, if they exist.
michael@0: 
michael@0:     Args:
michael@0:       flags: A list of flags to remove, eg. ['--single-process'].  Note that we
michael@0:              expect a complete match when removing flags; if you want to remove
michael@0:              a switch with a value, you must use the exact string used to add
michael@0:              it in the first place.
michael@0:     """
michael@0:     if flags:
michael@0:       assert flags[0] != 'chrome'
michael@0: 
michael@0:     for flag in flags:
michael@0:       if flag in self._current_flags:
michael@0:         self._current_flags.remove(flag)
michael@0:     self._UpdateCommandLineFile()
michael@0: 
michael@0:   def Restore(self):
michael@0:     """Restores the flags to their original state."""
michael@0:     self._current_flags = self._TokenizeFlags(self._orig_line)
michael@0:     self._UpdateCommandLineFile()
michael@0: 
michael@0:   def _UpdateCommandLineFile(self):
michael@0:     """Writes out the command line to the file, or removes it if empty."""
michael@0:     print "Current flags: ", self._current_flags
michael@0: 
michael@0:     if self._current_flags:
michael@0:       self._android_cmd.SetFileContents(CHROME_COMMAND_FILE,
michael@0:                                         'chrome ' +
michael@0:                                         ' '.join(self._current_flags))
michael@0:     else:
michael@0:       self._android_cmd.RunShellCommand('rm ' + CHROME_COMMAND_FILE)
michael@0: 
michael@0:   def _TokenizeFlags(self, line):
michael@0:     """Changes the string containing the command line into a list of flags.
michael@0: 
michael@0:     Follows similar logic to CommandLine.java::tokenizeQuotedArguments:
michael@0:     * Flags are split using whitespace, unless the whitespace is within a
michael@0:       pair of quotation marks.
michael@0:     * Unlike the Java version, we keep the quotation marks around switch
michael@0:       values since we need them to re-create the file when new flags are
michael@0:       appended.
michael@0: 
michael@0:     Args:
michael@0:       line: A string containing the entire command line.  The first token is
michael@0:             assumed to be the program name.
michael@0:     """
michael@0:     if not line:
michael@0:       return []
michael@0: 
michael@0:     tokenized_flags = []
michael@0:     current_flag = ""
michael@0:     within_quotations = False
michael@0: 
michael@0:     # Move through the string character by character and build up each flag
michael@0:     # along the way.
michael@0:     for c in line.strip():
michael@0:       if c is '"':
michael@0:         if len(current_flag) > 0 and current_flag[-1] == '\\':
michael@0:           # Last char was a backslash; pop it, and treat this " as a literal.
michael@0:           current_flag = current_flag[0:-1] + '"'
michael@0:         else:
michael@0:           within_quotations = not within_quotations
michael@0:           current_flag += c
michael@0:       elif not within_quotations and (c is ' ' or c is '\t'):
michael@0:         if current_flag is not "":
michael@0:           tokenized_flags.append(current_flag)
michael@0:           current_flag = ""
michael@0:       else:
michael@0:         current_flag += c
michael@0: 
michael@0:     # Tack on the last flag.
michael@0:     if not current_flag:
michael@0:       if within_quotations:
michael@0:         warnings.warn("Unterminated quoted string: " + current_flag)
michael@0:     else:
michael@0:       tokenized_flags.append(current_flag)
michael@0: 
michael@0:     # Return everything but the program name.
michael@0:     return tokenized_flags[1:]