From 4ed5c51cff98f852ee208a6856a7981d223a2253 Mon Sep 17 00:00:00 2001
From: Takeshi KOMIYA <i.tkomiya@gmail.com>
Date: Wed, 12 Sep 2018 20:09:22 +0900
Subject: [PATCH] logging: Add prefixed_warnings() helper

---
 doc/extdev/logging.rst     |  2 ++
 sphinx/util/logging.py     | 62 ++++++++++++++++++++++++++++++++++++++
 tests/test_util_logging.py | 21 ++++++++++++-
 3 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/doc/extdev/logging.rst b/doc/extdev/logging.rst
index 0d96c4eb9..005a2f87a 100644
--- a/doc/extdev/logging.rst
+++ b/doc/extdev/logging.rst
@@ -62,3 +62,5 @@ Logging API
 .. autofunction:: pending_logging()
 
 .. autofunction:: pending_warnings()
+
+.. autofunction:: prefixed_warnings()
diff --git a/sphinx/util/logging.py b/sphinx/util/logging.py
index a5bd7cd05..02bb12dd3 100644
--- a/sphinx/util/logging.py
+++ b/sphinx/util/logging.py
@@ -287,6 +287,53 @@ def skip_warningiserror(skip=True):
                 handler.removeFilter(disabler)
 
 
+@contextmanager
+def prefixed_warnings(prefix):
+    # type: (unicode) -> Generator
+    """Prepend prefix to all records for a while.
+
+    For example::
+
+        >>> with prefixed_warnings("prefix:"):
+        >>>     logger.warning('Warning message!')  # => prefix: Warning message!
+
+    .. versionadded:: 2.0
+    """
+    logger = logging.getLogger(NAMESPACE)
+    warning_handler = None
+    for handler in logger.handlers:
+        if isinstance(handler, WarningStreamHandler):
+            warning_handler = handler
+            break
+    else:
+        # warning stream not found
+        yield
+        return
+
+    prefix_filter = None
+    for _filter in warning_handler.filters:
+        if isinstance(_filter, MessagePrefixFilter):
+            prefix_filter = _filter
+            break
+
+    if prefix_filter:
+        # already prefixed
+        try:
+            previous = prefix_filter.prefix
+            prefix_filter.prefix = prefix
+            yield
+        finally:
+            prefix_filter.prefix = previous
+    else:
+        # not prefixed yet
+        try:
+            prefix_filter = MessagePrefixFilter(prefix)
+            warning_handler.addFilter(prefix_filter)
+            yield
+        finally:
+            warning_handler.removeFilter(prefix_filter)
+
+
 class LogCollector:
     def __init__(self):
         # type: () -> None
@@ -395,6 +442,21 @@ class DisableWarningIsErrorFilter(logging.Filter):
         return True
 
 
+class MessagePrefixFilter(logging.Filter):
+    """Prepend prefix to all records."""
+
+    def __init__(self, prefix):
+        # type: (unicode) -> None
+        self.prefix = prefix
+        super(MessagePrefixFilter, self).__init__()
+
+    def filter(self, record):
+        # type: (logging.LogRecord) -> bool
+        if self.prefix:
+            record.msg = self.prefix + ' ' + record.msg  # type: ignore
+        return True
+
+
 class SphinxLogRecordTranslator(logging.Filter):
     """Converts a log record to one Sphinx expects
 
diff --git a/tests/test_util_logging.py b/tests/test_util_logging.py
index 98affa886..fa7921cd1 100644
--- a/tests/test_util_logging.py
+++ b/tests/test_util_logging.py
@@ -20,7 +20,7 @@ from sphinx.errors import SphinxWarning
 from sphinx.testing.util import strip_escseq
 from sphinx.util import logging
 from sphinx.util.console import colorize
-from sphinx.util.logging import is_suppressed_warning
+from sphinx.util.logging import is_suppressed_warning, prefixed_warnings
 from sphinx.util.parallel import ParallelTasks
 
 
@@ -330,3 +330,22 @@ def test_skip_warningiserror(app, status, warning):
         with logging.pending_warnings():
             with logging.skip_warningiserror(False):
                 logger.warning('message')
+
+
+def test_prefixed_warnings(app, status, warning):
+    logging.setup(app, status, warning)
+    logger = logging.getLogger(__name__)
+
+    logger.warning('message1')
+    with prefixed_warnings('PREFIX:'):
+        logger.warning('message2')
+        with prefixed_warnings('Another PREFIX:'):
+            logger.warning('message3')
+        logger.warning('message4')
+    logger.warning('message5')
+
+    assert 'WARNING: message1' in warning.getvalue()
+    assert 'WARNING: PREFIX: message2' in warning.getvalue()
+    assert 'WARNING: Another PREFIX: message3' in warning.getvalue()
+    assert 'WARNING: PREFIX: message4' in warning.getvalue()
+    assert 'WARNING: message5' in warning.getvalue()