Better traceback with Config objects

PiperOrigin-RevId: 867552564

Author: Conchylicultor

kauldron/konfig/configdict_base.py

@@ -69,6 +69,9 @@ def __init__(
     init_dict = dict(init_dict or {})
     init_dict = _maybe_update_init_dict(init_dict)  # pytype: disable=name-error
 
+    # Capture the frame stack (to trace back where the ConfigDict is created)
+    object.__setattr__(self, '_frame', utils.FrameStack.from_current())
+
     # Normalize here rather than at the individul field level (`__setattr__`),
     # to have a global cache for all shared values (so shared fields are
     # correctly handled).

kauldron/konfig/configdict_proxy.py

@@ -128,6 +128,8 @@ def resolve(cfg, *, freeze=True):
   Returns:
     The resolved config.
   """
+  # Hide the function from the traceback (in Pytest and IPython 7).
+  __tracebackhide__ = True  # pylint: disable=unused-variable,invalid-name
 
   # Check if the config has a `_konfig_experimental_nofreeze` key.
   # TODO(klausg): make freeze=False the default and remove this.
@@ -140,6 +142,7 @@ def resolve(cfg, *, freeze=True):
     return _ConstructorResolver(freeze=freeze)._resolve_value(cfg)  # pylint: disable=protected-access
   except Exception as e:  # pylint: disable=broad-exception-caught
     logging.info(f'Full config (failing): {cfg}')  # pylint: disable=logging-fstring-interpolation
+    utils.filter_traceback(e.__traceback__)
     epy.reraise(e, 'Error resolving the config:\n')
 
 
@@ -250,8 +253,11 @@ def _resolve_dict(self, value):
         for k, v in kwargs.items()
     }
     args = [kwargs.pop(str(i)) for i in range(num_args(kwargs))]
-    with epy.maybe_reraise(prefix=lambda: _make_cfg_error_msg(value)):
+    try:
       obj = constructor(*args, **kwargs)
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      e = _wrap_cfg_error(e, value, frame=value._frame)  # pylint: disable=protected-access
+      raise e from e.__cause__
     # Allow the object to save the config it is comming from.
     if hasattr(type(obj), '__post_konfig_resolve__'):
       obj.__post_konfig_resolve__(value)
@@ -283,12 +289,40 @@ def num_args(obj: Mapping[str, Any]) -> int:
   return arg_id  # pylint: disable=undefined-loop-variable,undefined-variable
 
 
-def _make_cfg_error_msg(cfg: ml_collections.ConfigDict) -> str:
+def _wrap_cfg_error(
+    e: Exception,
+    cfg: ml_collections.ConfigDict,
+    frame: utils.FrameStack,
+) -> Exception:
+  """Wrap the exception with the config information."""
   cfg_str = repr(cfg)
   cfg_str = cfg_str.removeprefix('<ConfigDict[').removesuffix(']>')
   if len(cfg_str) > 300:  # `textwrap.shorten` remove `\n` so don't use it
     cfg_str = cfg_str[:295] + '[...]'
-  return f'Error while constructing cfg: {cfg_str}\n'
+  msg = (
+      f'For: {cfg_str}\n'
+      'Look at the the exception cause above to see where the ConfigDict'
+      ' was created.\n'
+      '======================================================================\n'
+  )
+
+  cause = ValueError('The ConfigDict was created here.').with_traceback(
+      frame.as_traceback()
+  )
+
+  # Backward compatibility.
+  if not hasattr(epy.reraise_utils, 'wrap_error'):
+    try:
+      epy.reraise(e, prefix=msg)
+    except Exception as exc:  # pylint: disable=broad-exception-caught
+      e = exc
+  else:
+    e = epy.reraise_utils.wrap_error(e, prefix=msg)
+
+  while e.__cause__ is not None:
+    e = e.__cause__
+  e.__cause__ = cause
+  return e
 
 
 def _as_dict(values: Mapping[str, Any]) -> dict[str, Any]:

kauldron/konfig/utils.py

@@ -14,11 +14,15 @@
 
 """Utils."""
 
+from __future__ import annotations
+
 import ast
 import dataclasses
 import functools
 import itertools
 import os
+import sys
+import types
 import typing
 from typing import Any, Generic, TypeVar
 
@@ -172,3 +176,56 @@ def maybe_decode_json_key(key: Any) -> Any:
     return ast.literal_eval(key.removeprefix(_JSON_RAW_PREFIX))
   else:
     return key
+
+
+@dataclasses.dataclass(slots=True, kw_only=True)
+class FrameInfo:
+  frame: types.FrameType
+  lasti: int
+  lineno: int
+
+
+class FrameStack(list[FrameInfo]):
+  """A list of FrameInfo objects.
+
+  Allow to reraise the stack trace from the current frame later on.
+
+  Note: This keep the frame objects alive.
+  """
+
+  @classmethod
+  def from_current(cls, depth: int = 0) -> FrameStack:
+    """Returns the current frame stack."""
+    stack = cls()
+    frame = sys._getframe(depth + 1)  # pylint: disable=protected-access
+    while frame:
+      # Do not capture internals IPython frames.
+      if frame.f_code.co_filename.endswith('IPython/core/interactiveshell.py'):
+        break
+      # Skip the konfig frames.
+      if '/kauldron/konfig/' in frame.f_code.co_filename:
+        frame = frame.f_back
+        continue
+
+      # Capture the current f_lasti and f_lineno as those values get updated
+      # as execution continues.
+      info = FrameInfo(frame=frame, lasti=frame.f_lasti, lineno=frame.f_lineno)  # pytype: disable=attribute-error
+      stack.append(info)
+      frame = frame.f_back  # pytype: disable=attribute-error
+    return stack
+
+  def as_traceback(self) -> types.TracebackType:
+    """Returns a traceback object for the current stack."""
+    tb = None
+    for info in self:
+      tb = types.TracebackType(tb, info.frame, info.lasti, info.lineno)
+    return tb  # pytype: disable=bad-return-type
+
+
+def filter_traceback(tb: types.TracebackType) -> None:
+  """Returns a filtered version of the traceback."""
+  # TODO(epot): Should also actually remove the tb for non-IPython frames.
+  while tb:
+    if '/kauldron/konfig/' in tb.tb_frame.f_code.co_filename:
+      tb.tb_frame.f_locals['__tracebackhide__'] = True
+    tb = tb.tb_next