python/mozperftest/mozperftest/environment.py

   1 # This Source Code Form is subject to the terms of the Mozilla Public
   2 # License, v. 2.0. If a copy of the MPL was not distributed with this
   3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   4 import contextlib
   5 import copy
   6
   7 from mozperftest.argparser import FLAVORS
   8 from mozperftest.hooks import Hooks
   9 from mozperftest.layers import Layers, StopRunError
  10 from mozperftest.metrics import pick_metrics
  11 from mozperftest.system import pick_system
  12 from mozperftest.test import pick_test
  13 from mozperftest.utils import MachLogger
  14
  15 SYSTEM, TEST, METRICS = 0, 1, 2
  16
  17
  18 class MachEnvironment(MachLogger):
  19     def __init__(self, mach_cmd, flavor="desktop-browser", hooks=None, **kwargs):
  20         MachLogger.__init__(self, mach_cmd)
  21         self._mach_cmd = mach_cmd
  22         self._mach_args = dict(
  23             [(self._normalize(key), value) for key, value in kwargs.items()]
  24         )
  25         self.layers = []
  26         if flavor not in FLAVORS:
  27             raise NotImplementedError(flavor)
  28         for layer in (pick_system, pick_test, pick_metrics):
  29             self.add_layer(layer(self, flavor, mach_cmd))
  30         if hooks is None:
  31             # we just load an empty Hooks instance
  32             hooks = Hooks(mach_cmd)
  33         self.hooks = hooks
  34
  35     @contextlib.contextmanager
  36     def frozen(self):
  37         self.freeze()
  38         try:
  39             # used to trigger __enter__/__exit__
  40             with self:
  41                 yield self
  42         finally:
  43             self.unfreeze()
  44
  45     def _normalize(self, name):
  46         if name.startswith("--"):
  47             name = name[2:]
  48         return name.replace("-", "_")
  49
  50     def set_arg(self, name, value):
  51         """Sets the argument"""
  52         # see if we want to restrict to existing keys
  53         self._mach_args[self._normalize(name)] = value
  54
  55     def get_arg(self, name, default=None, layer=None):
  56         name = self._normalize(name)
  57         marker = object()
  58         res = self._mach_args.get(name, marker)
  59         if res is marker:
  60             # trying with the name prefixed with the layer name
  61             if layer is not None and not name.startswith(layer.name):
  62                 name = "%s_%s" % (layer.name, name)
  63                 return self._mach_args.get(name, default)
  64             return default
  65         return res
  66
  67     def get_layer(self, name):
  68         for layer in self.layers:
  69             if isinstance(layer, Layers):
  70                 found = layer.get_layer(name)
  71                 if found is not None:
  72                     return found
  73             elif layer.name == name:
  74                 return layer
  75         return None
  76
  77     def add_layer(self, layer):
  78         self.layers.append(layer)
  79
  80     def freeze(self):
  81         # freeze args  (XXX do we need to freeze more?)
  82         self._saved_mach_args = copy.deepcopy(self._mach_args)
  83
  84     def unfreeze(self):
  85         self._mach_args = self._saved_mach_args
  86         self._saved_mach_args = None
  87
  88     def run(self, metadata):
  89         # run the system and test layers
  90         try:
  91             with self.layers[SYSTEM] as syslayer, self.layers[TEST] as testlayer:
  92                 metadata = testlayer(syslayer(metadata))
  93
  94             # then run the metrics layers
  95             with self.layers[METRICS] as metrics:
  96                 metadata = metrics(metadata)
  97         except StopRunError:
  98             # ends the cycle but without bubbling up the error
  99             pass
 100         return metadata
 101
 102     def __enter__(self):
 103         return self
 104
 105     def __exit__(self, type, value, traceback):
 106         return