Add `slave_alloc_update` tool.

scripts/common/slave_alloc.py

 # Copyright 2015 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.
 
 """Library to generate, maintain, and read static slave pool maps."""
 
 import collections
 import itertools
 import json
 import os
@@ skipping 56 matching lines @@
   external JSON file. This can be used to enforce class-to-slave mapping
   consistency (i.e., builder affinity).
 
   When a new allocation is performed, the SlaveAllocator's State is updated, and
   subsequent operations will prefer the previous layout.
   """
 
   # The default path to load/save state to, if none is specified.
   DEFAULT_STATE_PATH = 'slave_pool.json'
 
-  def __init__(self):
-    """Initializes a new slave pool instance."""
+  def __init__(self, state_path=None, list_unallocated=False):
+    """Initializes a new slave pool instance.
+
+    Args:
+      state_path (str): The path (relative or absolute) of the allocation
+          save-state JSON file. If None, DEFAULT_STATE_PATH will be used.
+      list_unallocated (bool): Include an entry listing unallocated slaves.
+          This entry will be ignored for operations, but can be useful when
+          generating expectations.
+    """
+    self._state_path = state_path
+    self._list_unallocated = list_unallocated
+
     self._state = None
     self._pools = {}
     self._classes = {}
     self._membership = {}
     self._all_slaves = {}
 
+  @property
+  def state_path(self):
+    return self._state_path or self.DEFAULT_STATE_PATH
+
   def LoadStateDict(self, state_class_map=None):
     """Loads previous allocation state from a state dictionary.
 
     The state dictionary is structured:
       <class-name>: {
         <class-subtype>: [
           <slave>
           ...
         ],
         ...
       }
 
     Args:
       state_class_map (dict): A state class map dictionary. If None or empty,
           the current state will be cleared.
     """
     if not state_class_map:
       self._state = None
       return
 
     class_map = {}
     for class_name, class_name_entry in state_class_map.iteritems():
       for subtype, slave_list in class_name_entry.iteritems():
         cls = SlaveClass(name=class_name, subtype=subtype)
         class_map.setdefault(cls, []).extend(str(s) for s in slave_list)
     self._state = SlaveState(
         class_map=class_map,
         unallocated=None)
 
-  def LoadState(self, path=None, enforce=True):
+  def LoadState(self, enforce=True):
     """Loads slave pools from the store, replacing the current in-memory set.
 
     Args:
-      path (str): If provided, the path to load from; otherwise,
-          DEFAULT_STATE_PATH will be used.
       enforce (bool): If True, raise an IOError if the state file does not
           exist or a ValueError if it could not be loaded.
     """
     state = {}
-    path = path or self.DEFAULT_STATE_PATH
-    if not os.path.exists(path):
+    if not os.path.exists(self.state_path):
       if enforce:
-        raise IOError("State path does not exist: %s" % (path,))
+        raise IOError("State path does not exist: %s" % (self.state_path,))
     try:
-      with open(path or self.DEFAULT_STATE_PATH, 'r') as fd:
+      with open(self.state_path, 'r') as fd:
         state = json.load(fd)
     except (IOError, ValueError):
       if enforce:
         raise
     self.LoadStateDict(state.get('class_map'))
 
-  def SaveState(self, path=None, list_unallocated=False):
-    """Saves the current slave pool set to the store path.
-
-    Args:
-      path (str): The path of the state file. If None, use DEFAULT_STATE_PATH.
-      list_unallocated (bool): Include an entry listing unallocated slaves.
-          This entry will be ignored for operations, but can be useful when
-          generating expectations.
-    """
+  def SaveState(self):
+    """Saves the current slave pool set to the store path."""
     state_dict = {}
     if self._state and self._state.class_map:
       class_map = state_dict['class_map'] = {}
       for sc, slave_list in self._state.class_map.iteritems():
         class_dict = class_map.setdefault(sc.name, {})
         subtype_dict = class_dict.setdefault(sc.subtype, [])
         subtype_dict.extend(slave_list)
 
-    if list_unallocated:
+    if self._list_unallocated:
       state_dict['unallocated'] = list(self._state.unallocated or ())
 
-    with open(path or self.DEFAULT_STATE_PATH, 'w') as fd:
+    with open(self.state_path, 'w') as fd:
       json.dump(state_dict, fd, sort_keys=True, indent=2)
 
   def AddPool(self, name, *slaves):
     """Returns (str): The slave pool that was allocated (for chaining).
 
     Args:
       name (str): The slave pool name.
       slaves: Slave name strings that belong to this pool.
     """
     pool = self._pools.get(name)
@@ skipping 167 matching lines @@
 
     # Convert SlaveMapEntry fields to immutable form.
     result = SlaveMap(
         entries={},
         unallocated=frozenset(n_state.unallocated))
     for k, v in slave_map_entries.iteritems():
       result.entries[k] = SlaveMapEntry(
           classes=frozenset(v.classes),
           keys=tuple(sorted(v.keys)))
     return result