Add LogDog / annotation protobuf support.

recipe_engine/stream.py

 # Copyright 2015 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """Abstract stream interface for representing recipe runs.
 
 We need to create streams for steps (and substeps) and also LOG_LINE steps.
 LogDog will implement LOG_LINE steps as real logs (i.e. uniformly), but
 annotations will implement them differently from normal logs, so we need
 a way to distinguish.
 
 StreamEngine will coordinate the multiplexing of streams.  In the case of
 annotations, this involves keeping track of the STEP_CURSOR and setting it
 accordingly, as well as filtering @@@ lines.
 
 Stream is a virtual well-behaved stream (associated with an Engine) which you
 can just write to without worrying.
 """
 
 import json
 import time
 
 from . import env
 from . import recipe_api
+from . import util
 
 
 class StreamEngine(object):
   class Stream(object):
     def write_line(self, line):
       raise NotImplementedError()
 
     def write_split(self, string):
       """Write a string (which may contain newlines) to the stream.  It will
       be terminated by a newline."""
@@ skipping 66 matching lines @@
     return self
 
   def __exit__(self, _exc_type, _exc_val, _exc_tb):
     self.close()
 
 
 # Because StreamEngine has no observations (i.e. it is an F-Algebra), we can
 # form products.  This code is entirely mechanical from the types (if we
 # had them formalized...).
 class ProductStreamEngine(StreamEngine):
-  def __init__(self, engine_a, engine_b):
-    assert engine_a and engine_b
-    self._engine_a = engine_a
-    self._engine_b = engine_b
+  def __init__(self, base, *engines):
+    self._engines = [base] + list(engines)
+    assert all(self._engines)
 
   class Stream(StreamEngine.Stream):
-    def __init__(self, stream_a, stream_b):
-      assert stream_a and stream_b
-      self._stream_a = stream_a
-      self._stream_b = stream_b
+    def __init__(self, *streams):
+      assert all(streams)
+      self._streams = streams
 
     def write_line(self, line):
-      self._stream_a.write_line(line)
-      self._stream_b.write_line(line)
+      for s in self._streams:
+        s.write_line(line)
 
     def close(self):
-      self._stream_a.close()
-      self._stream_b.close()
+      util.defer_exceptions_for(self._streams, lambda s: s.close())
 
   class StepStream(Stream):
     def _void_product(method_name):
       def inner(self, *args):
-        getattr(self._stream_a, method_name)(*args)
-        getattr(self._stream_b, method_name)(*args)
+        for s in self._streams:
+          getattr(s, method_name)(*args)
       return inner
 
     def new_log_stream(self, log_name):
       return ProductStreamEngine.Stream(
-          self._stream_a.new_log_stream(log_name),
-          self._stream_b.new_log_stream(log_name))
+          *(se.new_log_stream(log_name) for se in self._streams))
 
     add_step_text = _void_product('add_step_text')
     add_step_summary_text = _void_product('add_step_summary_text')
     add_step_link = _void_product('add_step_link')
     reset_subannotation_state = _void_product('reset_subannotation_state')
     set_step_status = _void_product('set_step_status')
     set_build_property = _void_product('set_build_property')
     trigger = _void_product('trigger')
 
   def new_step_stream(self, step_config):
     return self.StepStream(
-        self._engine_a.new_step_stream(step_config),
-        self._engine_b.new_step_stream(step_config))
+        *(se.new_step_stream(step_config)
+          for se in self._engines))
 
   def open(self):
-    self._engine_a.open()
-    self._engine_b.open()
+    for se in self._engines:
+      se.open()
 
   def close(self):
-    self._engine_a.close()
-    self._engine_b.close()
+    util.defer_exceptions_for(self._engines, lambda se: se.close())
+
+  def append_stream_engine(self, se):
+    assert isinstance(se, StreamEngine)
+    self._engines.append(se)
 
 
 def _noop(*args, **kwargs):
   pass
 
 class NoopStreamEngine(StreamEngine):
   class Stream(StreamEngine.Stream):
     write_line = _noop
     close = _noop
 
@@ skipping 246 matching lines @@
   string (which means it might not be valid ascii), we decode the string with
   the 'replace' error mode, which replaces invalid characters with a suitable
   replacement character.
   """
   try:
     return str(s)
   except UnicodeEncodeError:
     return s.encode('utf-8', 'replace')
   except UnicodeDecodeError:
     return s.decode('utf-8', 'replace')