From 9c0086afd0220248732973ed485749f3a70fde46 Mon Sep 17 00:00:00 2001
From: Daiyi Peng <daiyip@google.com>
Date: Sat, 26 Apr 2025 16:03:43 -0700
Subject: [PATCH] Add `pg.logging.register_frame_to_skip` to better support
 logger wrappers.

PiperOrigin-RevId: 751851424
---
 pyglove/core/logging.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/pyglove/core/logging.py b/pyglove/core/logging.py
index 91441b06..f9656147 100644
--- a/pyglove/core/logging.py
+++ b/pyglove/core/logging.py
@@ -17,17 +17,47 @@
 events without introducing library dependencies in PyGlove.
 """
 
+import inspect
 import logging
+from typing import Any, Callable, List, Union
 
 
 _DEFAULT_LOGGER = logging.getLogger()
 
 
+def register_frame_to_skip(
+    method: Union[Callable[..., Any], List[Callable[..., Any]]]
+) -> bool:
+  """Skips the source of the given method when logging.
+
+  Args:
+    method: The method to skip. Can be a single method or a list of methods.
+
+  Returns:
+    True if the method is registered to skip.
+
+  Raises:
+    TypeError: The source file of the method cannot be inspected.
+  """
+  register_fn = getattr(
+      _DEFAULT_LOGGER.__class__, 'register_frame_to_skip', None
+  )
+  if register_fn is None:
+    return False
+  methods = [method] if not isinstance(method, list) else method
+  for m in methods:
+    register_fn(inspect.getsourcefile(m), m.__name__)
+  return True
+
+
 def set_logger(logger: logging.Logger) -> None:
   """Sets current logger."""
   global _DEFAULT_LOGGER
   _DEFAULT_LOGGER = logger
 
+  # Skip logging frames in pyglove.logging.
+  register_frame_to_skip([debug, info, warning, error, critical])
+
 
 def get_logger() -> logging.Logger:
   """Gets the current logger."""