From b3bfa346c36150cabd123e9998ce03c65fda7630 Mon Sep 17 00:00:00 2001
From: Alex Ford <a.sewall.ford@gmail.com>
Date: Mon, 22 Mar 2021 06:15:16 -0700
Subject: [PATCH] Using existing logging level if defined in addLoggingLevel
 (#5632)

* Use existing level if defined in addLoggingLevel

Fixes https://github.com/iterative/dvc/issues/4546

* Add addLoggingLevel test

* Format fix.
---
 dvc/logger.py             | 28 +++++++++++++++++++++-------
 tests/unit/test_logger.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/dvc/logger.py b/dvc/logger.py
index d48dc3437..396fb242e 100644
--- a/dvc/logger.py
+++ b/dvc/logger.py
@@ -23,14 +23,19 @@ def addLoggingLevel(levelName, levelNum, methodName=None):
     Adds a new logging level to the `logging` module and the
     currently configured logging class.
 
+    Uses the existing numeric levelNum if already defined.
+
     Based on https://stackoverflow.com/questions/2183233
     """
     if methodName is None:
         methodName = levelName.lower()
 
-    assert not hasattr(logging, levelName)
-    assert not hasattr(logging, methodName)
-    assert not hasattr(logging.getLoggerClass(), methodName)
+    # If the level name is already defined as a top-level `logging`
+    # constant, then adopt the existing numeric level.
+    if hasattr(logging, levelName):
+        existingLevelNum = getattr(logging, levelName)
+        assert isinstance(existingLevelNum, int)
+        levelNum = existingLevelNum
 
     def logForLevel(self, message, *args, **kwargs):
         if self.isEnabledFor(levelNum):
@@ -40,10 +45,19 @@ def addLoggingLevel(levelName, levelNum, methodName=None):
     def logToRoot(message, *args, **kwargs):
         logging.log(levelNum, message, *args, **kwargs)
 
-    logging.addLevelName(levelNum, levelName)
-    setattr(logging, levelName, levelNum)
-    setattr(logging.getLoggerClass(), methodName, logForLevel)
-    setattr(logging, methodName, logToRoot)
+    # getLevelName resolves the numeric log level if already defined,
+    # otherwise returns a string
+    if not isinstance(logging.getLevelName(levelName), int):
+        logging.addLevelName(levelNum, levelName)
+
+    if not hasattr(logging, levelName):
+        setattr(logging, levelName, levelNum)
+
+    if not hasattr(logging.getLoggerClass(), methodName):
+        setattr(logging.getLoggerClass(), methodName, logForLevel)
+
+    if not hasattr(logging, methodName):
+        setattr(logging, methodName, logToRoot)
 
 
 class LoggingException(Exception):
diff --git a/tests/unit/test_logger.py b/tests/unit/test_logger.py
index 37a725e68..dc766471b 100644
--- a/tests/unit/test_logger.py
+++ b/tests/unit/test_logger.py
@@ -248,3 +248,31 @@ def test_info_with_debug_loglevel_shows_no_datetime(caplog, dt):
         logger.info("message")
 
         assert "message" == formatter.format(caplog.records[0])
+
+
+def test_add_existing_level(caplog, dt):
+    # Common pattern to configure logging level in external libraries
+    # eg:
+    # https://github.com/bokeh/bokeh/blob/04bb30fef2e72e64baaa8b2f330806d5bfdd3b11/
+    # bokeh/util/logconfig.py#L79-L85
+    TRACE2 = 4
+    logging.addLevelName(TRACE2, "TRACE2")
+    logging.TRACE2 = TRACE2
+
+    dvc.logger.addLoggingLevel("TRACE2", 2)
+
+    # DVC sets all expected entrypoints, but doesn't override the level
+    assert logging.TRACE2 == 4
+    assert hasattr(logging, "trace2")
+    assert hasattr(logger, "trace2")
+    assert logging.getLevelName("TRACE2") == 4
+
+    # The TRACE2 logging level uses the original, higher logging level
+    with caplog.at_level(logging.TRACE2, logger="dvc"):
+        logger.trace2("TRACE2")
+    assert len(caplog.records) == 1
+
+    (record,) = caplog.records
+    assert record.levelno == 4
+    assert record.levelname == "TRACE2"
+    assert record.message == "TRACE2"