Augmentations: Allow configuring of augmentations per workload

Add support for enabling and disabling of augmentations (instruments and
output processors) on a per workload basis.

wa/framework/output_processor.py

@@ -5,6 +5,7 @@ from wa.framework.exception import ConfigError
 from wa.framework.instruments import is_installed
 from wa.framework.plugin import Plugin
 from wa.utils.log import log_error, indent, dedent
+from wa.utils.misc import isiterable
 
 
 class OutputProcessor(Plugin):
@@ -12,6 +13,10 @@ class OutputProcessor(Plugin):
     kind = 'output_processor'
     requires = []
 
+    def __init__(self, **kwargs):
+        super(OutputProcessor, self).__init__(**kwargs)
+        self.is_enabled = True
+
     def validate(self):
         super(OutputProcessor, self).validate()
         for instrument in self.requires:
@@ -40,6 +45,43 @@ class ProcessorManager(object):
         processor.logger.context = context
         self.processors.append(processor)
 
+    def disable_all(self):
+        for output_processor in self.processors:
+            self._disable_output_processor(output_processor)
+
+    def enable_all(self):
+        for output_processor in self.processors:
+            self._enable_output_processor(output_processor)
+
+    def enable(self, to_enable):
+        if isiterable(to_enable):
+            for inst in to_enable:
+                self._enable_output_processor(inst)
+        else:
+            self._enable_output_processor(to_enable)
+
+    def disable(self, to_disable):
+        if isiterable(to_disable):
+            for inst in to_disable:
+                self._disable_output_processor(inst)
+        else:
+            self._disable_output_processor(to_disable)
+
+    def get_output_processor(self, processor):
+        if isinstance(processor, OutputProcessor):
+            return processor
+
+        for p in self.processors:
+            if processor == p.name:
+                return p
+        raise ValueError('Output processor {} is not installed'.format(processor))
+
+    def get_enabled(self):
+        return [p for p in self.processors if p.is_enabled]
+
+    def get_disabled(self):
+        return [p for p in self.processors if not p.is_enabled]
+
     def validate(self):
         for proc in self.processors:
             proc.validate()
@@ -74,15 +116,29 @@ class ProcessorManager(object):
         try:
             indent()
             for proc in self.processors:
-                proc_func = getattr(proc, method_name, None)
-                if proc_func is None:
-                    continue
-                try:
-                    self.logger.info(message.format(proc.name))
-                    proc_func(*args)
-                except Exception as e:
-                    if isinstance(e, KeyboardInterrupt):
-                        raise
-                    log_error(e, self.logger)
+                if proc.is_enabled:
+                    proc_func = getattr(proc, method_name, None)
+                    if proc_func is None:
+                        continue
+                    try:
+                        self.logger.info(message.format(proc.name))
+                        proc_func(*args)
+                    except Exception as e:
+                        if isinstance(e, KeyboardInterrupt):
+                            raise
+                        log_error(e, self.logger)
         finally:
             dedent()
+
+    def _enable_output_processor(self, inst):
+        inst = self.get_output_processor(inst)
+        self.logger.debug('Enabling output processor {}'.format(inst.name))
+        if not inst.is_enabled:
+            inst.is_enabled = True
+
+    def _disable_output_processor(self, inst):
+        inst = self.get_output_processor(inst)
+        self.logger.debug('Disabling output processor {}'.format(inst.name))
+        if inst.is_enabled:
+            inst.is_enabled = False
+