import gzip
import bz2


def open_file(fname, mode):
  """ A wrapper around `open()` that also supports gzip and bz2.

  Args:
    mode (str): should either be 'r' or 'w'. A 't' will be appended when
                reading from a compressed file.
  """
  if fname.endswith('.gz'):
    return gzip.open(fname, mode + 't')
  elif fname.endswith('.bz2'):
    return bz2.open(fname, mode + 't')
  else:
    return open(fname, mode)