heron/common/src/python/utils/proc.py

# Copyright 2017 Twitter. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
''' proc.py: subprocess and subprocess's stdout/stderr management '''
from threading import Thread

def _stream_process_fileno(fileno, handler):
  """ Stream handling each line from fileno
  :param filno: file object
  :param handler: a function that will be called for each line from fileno
  :return: None
  """
  while 1:
    line = fileno.readline()
    if not line:
      break
    handler(line)

def stream_process_stdout(process, handler):
  """ Stream the stdout for a process out to display
  :param process: the process to stream the stdout for
  :param handler: a function that will be called for each stdout line
  :return: None
  """
  _stream_process_fileno(process.stdout, handler)

def stream_process_stderr(process, handler):
  """ Stream the stderr for a process out to display
  :param process: the process to stream the stderr for
  :param handler: a function that will be called for each stderr line
  :return: None
  """
  _stream_process_fileno(process.stderr, handler)

def _async_stream_process_output(process, stream_fn, handler):
  """ Stream and handle the output of a process
  :param process: the process to stream the output for
  :param stream_fn: the function that applies handler to process
  :param handler: a function that will be called for each log line
  :return: None
  """
  logging_thread = Thread(target=stream_fn, args=(process, handler, ))

  # Setting the logging thread as a daemon thread will allow it to exit with the program
  # rather than blocking the exit waiting for it to be handled manually.
  logging_thread.daemon = True
  logging_thread.start()

  return logging_thread

def async_stream_process_stdout(process, handler):
  """ Stream and handler the stdout of a process
  :param process: the process to stream the stdout for
  :param handler: a function that will be called to handle each line
  :return: None
  """
  return _async_stream_process_output(process, stream_process_stdout, handler)

def async_stream_process_stderr(process, handler):
  """ Stream and handler the stderr of a process
  :param process: the process to stream the stderr for
  :param handler: a function that will be called to handle each line
  :return: None
  """
  return _async_stream_process_output(process, stream_process_stderr, handler)

class StringBuilder(object):
  def __init__(self):
    self.str = ""

  def add(self, line):
    self.str += line

  def result(self):
    return self.str

def async_stdout_builder(proc):
  """ Save stdout into string builder
  :param proc: the process to save stdout for
  :return StringBuilder
  """
  stdout_builder = StringBuilder()
  async_stream_process_stdout(proc, stdout_builder.add)
  return stdout_builder

def async_stderr_builder(proc):
  """ Save stderr into string builder
  :param proc: the process to save stderr for
  :return StringBuilder
  """
  stderr_builder = StringBuilder()
  async_stream_process_stderr(proc, stderr_builder.add)
  return stderr_builder

def async_stdout_stderr_builder(proc):
  """ Save stdout and stderr into string builders
  :param proc: the process to save stdout and stderr for
  :return (StringBuilder, StringBuilder)
  """
  return async_stdout_builder(proc), async_stderr_builder(proc)