feat: Attribute etqdm log messages to the caller's source location.

etils/etqdm/tqdm_utils.py

@@ -14,6 +14,8 @@
 
 """Wrapper for tdqm."""
 
+import logging as logging_stdlib
+import sys
 import typing
 from typing import Optional, TypeVar
 
@@ -30,15 +32,46 @@
 
 
 class _LogFile:
-  """A File-like object that log to INFO."""
+  """A tqdm-compatible file-like object that logs to INFO.
 
-  def write(self, message):
-    logging.info(message)
+  Captures the caller's source location at construction time so that
+  log messages attribute to the code that created the progress bar,
+  not to tqdm internals.
+  """
 
-  def flush(self):
+  def __init__(self, caller_depth: int = 2) -> None:
+    """Initializes the log file.
+
+    Args:
+      caller_depth: Number of frames to skip above ``_LogFile.__init__`` to
+        reach the "real" caller.  Default is 2, which skips ``__init__`` itself
+        and one wrapper (e.g. ``tqdm()``).
+    """
+    frame = sys._getframe(caller_depth)
+    self._caller_file = frame.f_code.co_filename
+    self._caller_lineno = frame.f_lineno
+    self._caller_func = frame.f_code.co_name
+
+  def write(self, message: str) -> None:
+    """Logs a non-empty message at INFO level with the captured source location."""
+    if message := message.strip():
+      logger = logging.get_absl_logger()
+      record = logger.makeRecord(
+          name=logger.name,
+          level=logging_stdlib.INFO,
+          fn=self._caller_file,
+          lno=self._caller_lineno,
+          msg=message,
+          args=(),
+          exc_info=None,
+          func=self._caller_func,
+      )
+      logger.handle(record)
+
+  def flush(self) -> None:
     pass
 
-  def close(self):
+  def close(self) -> None:
     pass
 
 

etils/etqdm/tqdm_utils_test.py

@@ -12,10 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensorflow_datasets.core.utils.tqdm_utils."""
+import logging
+import logging.handlers
+import os
+import unittest
+from unittest import mock
 
 from etils import etqdm
+from etils.etqdm import tqdm_utils
 
 
-def test_disable_tqdm():
-  assert list(etqdm.tqdm(range(3))) == [0, 1, 2]
+class TqdmBasicTest(unittest.TestCase):
+  """Baseline test that tqdm wraps iterables correctly."""
+
+  def test_tqdm_iterates(self):
+    self.assertEqual(list(etqdm.tqdm(range(3))), [0, 1, 2])
+
+
+if __name__ == '__main__':
+  unittest.main()