""" Code to help using the logging module with tqdm progress bars. """

import contextlib
import logging

import threading


logging_handlers_lock = threading.Lock()


class TqdmLoggingHandler(logging.Handler):

  """ Logging handler sending messages to the tqdm write method (avoids overlap). """

  def __init__(self, tqdm, *args, **kwargs):
    self.tqdm = tqdm
    super().__init__(*args, **kwargs)

  def emit(self, record):
    msg = self.format(record)
    self.tqdm.write(msg)


@contextlib.contextmanager
def redirect_logging(tqdm_obj, logger=logging.getLogger()):
  """ Context manager to redirect logging to a TqdmLoggingHandler object and then restore the original logging behavior. """
  with logging_handlers_lock:
    # remove current handlers
    prev_handlers = []
    for handler in logger.handlers.copy():
      prev_handlers.append(handler)
      logger.removeHandler(handler)

    # add tqdm handler
    tqdm_handler = TqdmLoggingHandler(tqdm_obj)
    if prev_handlers[-1].formatter is not None:
      tqdm_handler.setFormatter(prev_handlers[-1].formatter)
    logger.addHandler(tqdm_handler)

  try:
    yield
  finally:
    # restore handlers
    with logging_handlers_lock:
      logger.removeHandler(tqdm_handler)
      for handler in prev_handlers:
        logger.addHandler(handler)
