1# Copyright 2017 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4"""Logging util functions.
5
6It would be named logging, but other modules in this directory use the default
7logging module, so that would break them.
8"""
9
10import contextlib
11import logging
12
13@contextlib.contextmanager
14def CaptureLogs(file_stream):
15  if not file_stream:
16    # No file stream given, just don't capture logs.
17    yield
18    return
19
20  fh = logging.StreamHandler(file_stream)
21
22  logger = logging.getLogger()
23  # Try to copy the current log format, if one is set.
24  if logger.handlers and hasattr(logger.handlers[0], 'formatter'):
25    fh.formatter = logger.handlers[0].formatter
26  else:
27    fh.setFormatter(logging.Formatter(
28        '(%(levelname)s) %(asctime)s %(message)s'))
29  logger.addHandler(fh)
30
31  try:
32    yield
33  finally:
34    logger = logging.getLogger()
35    logger.removeHandler(fh)
36