import six
import logging


def captured_logging(logger):
    """
    Utility for capturing a logger object's output.
    Useful for command line output testing.
    """

    class LogWrapper(object):
        def __init__(self):
            self.buffer = six.StringIO()
            self.stream_handler = logging.StreamHandler(self.buffer)
            self._output = None

        @property
        def output(self):
            if self._output is None:
                self.stream_handler.flush()
                # Standardize line endings
                self._output = self.buffer\
                    .getvalue().replace('\r\n', '\n')
            return self._output

    log_wrapper = LogWrapper()
    logger.addHandler(log_wrapper.stream_handler)
    yield log_wrapper
    logger.removeHandler(log_wrapper.stream_handler)