Class OutputStream

shell/impala_shell/shell_output.py:152–191

class OutputStream(object):
  def __init__(self, formatter, filename=None):
    """Helper class for writing query output.

    User should invoke the `write(data)` method of this object.
    `data` is a list of lists.
    """
    self.formatter = formatter
    self.filename = filename

  def write(self, data):
    formatted_data = self.formatter.format(data)
    if self.filename is not None:
      try:
        with open(self.filename, 'ab') as out_file:
          # Note that instances of this class do not persist, so it's fine to
          # close the we close the file handle after each write.
          # The file is opened in binary mode. Python 2 returns Unicode bytes
          # that can be written directly. Python 3 returns a string, which
          # we need to encode before writing.
          # TODO: Reexamine the contract of the format() function and see if
          # we can remove this.
          if sys.version_info.major == 2 and isinstance(formatted_data, str):
            out_file.write(formatted_data)
          else:
            out_file.write(formatted_data.encode('utf-8'))
          out_file.write(b'\n')
      except IOError as err:
        file_err_msg = "Error opening file %s: %s" % (self.filename, str(err))
        print('{0} (falling back to stderr)'.format(file_err_msg), file=sys.stderr)
        print(formatted_data, file=sys.stderr)
    else:
      # If filename is None, then just print to stdout
      print(formatted_data)

  def flush(self):
    # When outputing to a file, the file is currently closed with each write,
    # so the flush doesn't need to do anything.
    if self.filename is None:
      sys.stdout.flush()


class OverwritingStdErrOutputStream(object):