Add logger level control flag, take 2.

This adds --logger_levels, a flag that allows setting the log levels of
loggers by name.

This makes it easy to adjust log levels on a per-logger basis without having
to write application code to do so.

Changes from original:
 * Use empty dict as the default value to work around some dependencies
   relying on the default matching the parsed value.
 * Return a copy instead of `MappingProxyType` because values are expected
   to be compatible with `copy.deepcopy()`

PiperOrigin-RevId: 305493905
Change-Id: Ib826653b58d621aebc89aa1ded69647c534cde13

2 files changed, 99 insertions, 0 deletions

diff --git a/absl/logging/tests/logging_functional_test.py b/absl/logging/tests/logging_functional_test.py
index dd7e711..b8c79d3 100755
--- a/absl/logging/tests/logging_functional_test.py
+++ b/absl/logging/tests/logging_functional_test.py
@@ -647,6 +647,31 @@ E0000 00:00:00.000000 12345 logging_functional_test_helper.py:123] std error log
         test_name='bad_exc_info',
         use_absl_log_file=True)
 
+  def test_verbosity_logger_levels_flag_ordering(self):
+    """Make sure last-specified flag wins."""
+
+    def assert_error_level_logged(stderr):
+      lines = stderr.splitlines()
+      for line in lines:
+        self.assertIn('std error log', line)
+
+    self._exec_test(
+        _verify_ok,
+        test_name='std_logging',
+        expected_logs=[('stderr', None, assert_error_level_logged)],
+        extra_args=['-v=1', '--logger_levels=:ERROR'])
+
+    def assert_debug_level_logged(stderr):
+      lines = stderr.splitlines()
+      for line in lines:
+        self.assertRegex(line, 'std (debug|info|warning|error) log')
+
+    self._exec_test(
+        _verify_ok,
+        test_name='std_logging',
+        expected_logs=[('stderr', None, assert_debug_level_logged)],
+        extra_args=['--logger_levels=:ERROR', '-v=1'])
+
   def test_none_exc_info_py_logging(self):
 
     if six.PY2:
diff --git a/absl/logging/tests/logging_test.py b/absl/logging/tests/logging_test.py
index ddae7a1..290b2f3 100644
--- a/absl/logging/tests/logging_test.py
+++ b/absl/logging/tests/logging_test.py
@@ -59,6 +59,80 @@ class ConfigurationTest(absltest.TestCase):
                    logging.PythonFormatter))
 
 
+class LoggerLevelsTest(parameterized.TestCase):
+
+  def setUp(self):
+    super(LoggerLevelsTest, self).setUp()
+    # Since these tests muck with the flag, always save/restore in case the
+    # tests forget to clean up properly.
+    # enter_context() is py3-only, but manually enter/exit should suffice.
+    cm = self.set_logger_levels({})
+    cm.__enter__()
+    self.addCleanup(lambda: cm.__exit__(None, None, None))
+
+  @contextlib.contextmanager
+  def set_logger_levels(self, levels):
+    original_levels = {
+        name: std_logging.getLogger(name).level for name in levels
+    }
+
+    try:
+      with flagsaver.flagsaver(logger_levels=levels):
+        yield
+    finally:
+      for name, level in original_levels.items():
+        std_logging.getLogger(name).setLevel(level)
+
+  def assert_logger_level(self, name, expected_level):
+    logger = std_logging.getLogger(name)
+    self.assertEqual(logger.level, expected_level)
+
+  def assert_logged(self, logger_name, expected_msgs):
+    logger = std_logging.getLogger(logger_name)
+    # NOTE: assertLogs() sets the logger to INFO if not specified.
+    with self.assertLogs(logger, logger.level) as cm:
+      logger.debug('debug')
+      logger.info('info')
+      logger.warning('warning')
+      logger.error('error')
+      logger.critical('critical')
+
+    actual = {r.getMessage() for r in cm.records}
+    self.assertEqual(set(expected_msgs), actual)
+
+  @unittest.skipIf(six.PY2, 'Py2 is missing assertLogs')
+  def test_setting_levels(self):
+    # Other tests change the root logging level, so we can't
+    # assume it's the default.
+    orig_root_level = std_logging.root.getEffectiveLevel()
+    with self.set_logger_levels({'foo': 'ERROR', 'bar': 'DEBUG'}):
+
+      self.assert_logger_level('foo', std_logging.ERROR)
+      self.assert_logger_level('bar', std_logging.DEBUG)
+      self.assert_logger_level('', orig_root_level)
+
+      self.assert_logged('foo', {'error', 'critical'})
+      self.assert_logged('bar',
+                         {'debug', 'info', 'warning', 'error', 'critical'})
+
+  @parameterized.named_parameters(
+      ('empty', ''),
+      ('one_value', 'one:INFO'),
+      ('two_values', 'one.a:INFO,two.b:ERROR'),
+      ('whitespace_ignored', ' one : DEBUG , two : INFO'),
+  )
+  def test_serialize_parse(self, levels_str):
+    fl = FLAGS['logger_levels']
+    fl.parse(levels_str)
+    expected = levels_str.replace(' ', '')
+    actual = fl.serialize()
+    self.assertEqual('--logger_levels={}'.format(expected), actual)
+
+  def test_invalid_value(self):
+    with self.assertRaisesRegex(ValueError, 'Unknown level.*10'):
+      FLAGS['logger_levels'].parse('foo:10')
+
+
 class PythonHandlerTest(absltest.TestCase):
   """Tests the PythonHandler class."""