def capture_cpp_log_context()

in tensorflow_decision_forests/tensorflow/tf_logging.py [0:0]

• 45 lines of code
• 13 McCabe index (conditional complexity)

def capture_cpp_log_context(verbose=False):
  """Creates a context to display or hide the c++ code logs to the user.

  Tested with python, ipython, colab and jupyter notebook.
  In the internal build, only impact python's print.

  Args:
    verbose: If true, the training logs are shown in logging and python's print.
      If false, the training logs are not shown in logging nor python's print.

  Returns:
    A context.
  """

  # Does nothing
  @contextmanager
  def no_op_context():
    yield

  # Hide the Yggdrasil training logs.
  @contextmanager
  def hide_cpp_logs():
    # Stop displaying c++ logs.
    set_yggdrasil_logging_level(0)
    try:
      yield
    finally:
      # Re-start displaying c++ logs.
      set_yggdrasil_logging_level(2)

  if not verbose:
    # Make sure the c++ logs are not visible to the user.
    return hide_cpp_logs()

  def is_direct_output(stream):
    """Checks if output stream redirects to the shell/console directly."""

    if stream.isatty():
      return True
    if isinstance(stream, io.TextIOWrapper):
      return is_direct_output(stream.buffer)
    if isinstance(stream, io.BufferedWriter):
      return is_direct_output(stream.raw)
    if isinstance(stream, io.FileIO):
      return stream.fileno() in [1, 2]
    return False

  if ((REDIRECT_YGGDRASIL_CPP_OUTPUT_TO_PYTHON_OUTPUT == "auto" and
       is_direct_output(sys.stdout)) or
      not REDIRECT_YGGDRASIL_CPP_OUTPUT_TO_PYTHON_OUTPUT):
    # The cout and cerr of the c++ library are already visible to the user.
    return no_op_context()

  global REDIRECT_MESSAGE_WAS_PRINTED
  if not REDIRECT_MESSAGE_WAS_PRINTED:
    REDIRECT_MESSAGE_WAS_PRINTED = True
    info("Standard output detected as not visible to the user e.g. running "
         "in a notebook. Creating a training log redirection. If training get "
         "stuck, try calling tfdf.keras.set_training_logs_redirection(False).")

  # pytype: disable=import-error
  # pylint: disable=g-import-not-at-top
  # pylint: disable=g-importing-member
  # pylint: disable=bare-except

  # The cout and cerr of the c++ library are not visible to the user.
  # Redirect them to python's standard output.
  try:
    from colabtools.googlelog import CaptureLog
    return CaptureLog()
  except:
    try:
      from wurlitzer import sys_pipes
      # This can hang if the cout/cerr is visible to the user.
      return sys_pipes()
    except:
      warning("Cannot redirect the training output because neither of "
              "colabtools.googlelog or wurlitzer available. Run `pip install "
              "wurlitzer -U` and try again.")
      return no_op_context()