Drop check_shapes from except tracebacks.

Author: jesnie

check_shapes/__init__.py

@@ -22,9 +22,11 @@
     DocstringFormat,
     ShapeCheckingState,
     disable_check_shapes,
+    get_drop_frames,
     get_enable_check_shapes,
     get_enable_function_call_precompute,
     get_rewrite_docstrings,
+    set_drop_frames,
     set_enable_check_shapes,
     set_enable_function_call_precompute,
     set_rewrite_docstrings,
@@ -64,6 +66,7 @@
     "error_contexts",
     "exceptions",
     "get_check_shapes",
+    "get_drop_frames",
     "get_enable_check_shapes",
     "get_enable_function_call_precompute",
     "get_rewrite_docstrings",
@@ -73,6 +76,7 @@
     "inheritance",
     "parser",
     "register_get_shape",
+    "set_drop_frames",
     "set_enable_check_shapes",
     "set_enable_function_call_precompute",
     "set_rewrite_docstrings",

check_shapes/config.py

@@ -16,21 +16,11 @@
 """
 from contextlib import contextmanager
 from enum import Enum
-from typing import Callable, Iterator, List, Union
+from typing import Callable, ContextManager, Iterator, List, Union
 
-IsCompiledMode = Callable[[], bool]
-
-
-_is_compiled_mode: List[IsCompiledMode] = []
-
-
-def add_is_compiled_mode(is_compiled_mode: IsCompiledMode) -> None:
-    """
-    Add a function for determining whether we are currently executing "compiled mode".
+from dropstackframe import get_enable_drop_stack_frame, set_enable_drop_stack_frame
 
-    Used when func:`set_enable_check_shapes` is set to ``EAGER_MODE_ONLY``.
-    """
-    _is_compiled_mode.append(is_compiled_mode)
+IsCompiledMode = Callable[[], bool]
 
 
 class ShapeCheckingState(Enum):
@@ -67,9 +57,6 @@ def __bool__(self) -> bool:
             return False
 
 
-_enabled = ShapeCheckingState.EAGER_MODE_ONLY
-
-
 class DocstringFormat(Enum):
     """
     Enumeration of supported formats of docstrings.
@@ -86,11 +73,21 @@ class DocstringFormat(Enum):
     """
 
 
+_is_compiled_mode: List[IsCompiledMode] = []
+_enabled = ShapeCheckingState.EAGER_MODE_ONLY
 _docstring_format = DocstringFormat.SPHINX
-
 _function_call_precompute_enabled = False
 
 
+def add_is_compiled_mode(is_compiled_mode: IsCompiledMode) -> None:
+    """
+    Add a function for determining whether we are currently executing "compiled mode".
+
+    Used when func:`set_enable_check_shapes` is set to ``EAGER_MODE_ONLY``.
+    """
+    _is_compiled_mode.append(is_compiled_mode)
+
+
 def set_enable_check_shapes(enabled: Union[ShapeCheckingState, str, bool]) -> None:
     """
     Set whether to enable :mod:`check_shapes`.
@@ -192,3 +189,17 @@ def get_enable_function_call_precompute() -> bool:
     Get whether to precompute function call path and line numbers for debugging.
     """
     return _function_call_precompute_enabled
+
+
+def set_drop_frames(drop_frames: bool) -> ContextManager[None]:
+    """
+    Set whether :mod:`check_shapes` should hide itself from exception stack traces.
+    """
+    return set_enable_drop_stack_frame(drop_frames)
+
+
+def get_drop_frames() -> bool:
+    """
+    Get whether :mod:`check_shapes` should hide itself from exception stack traces.
+    """
+    return get_enable_drop_stack_frame()

check_shapes/decorator.py

@@ -19,6 +19,8 @@
 from functools import update_wrapper
 from typing import Any, Callable, Sequence, cast
 
+from dropstackframe import drop_stack_frame
+
 from .accessors import set_check_shapes
 from .argument_ref import RESULT_TOKEN
 from .base_types import C
@@ -117,7 +119,11 @@ def _check_shapes(func: C) -> C:
 
         def wrapped_function(*args: Any, **kwargs: Any) -> Any:
             if not get_enable_check_shapes():
-                return func(*args, **kwargs)
+                try:
+                    return func(*args, **kwargs)
+                except Exception:
+                    drop_stack_frame()
+                    raise
 
             try:
                 bound_arguments = signature.bind(*args, **kwargs)
@@ -182,7 +188,11 @@ def _check_specs(specs: Sequence[ParsedArgumentSpec]) -> None:
             _check_specs(pre_specs)
 
             with set_shape_checker(checker):
-                result = func(*args, **kwargs)
+                try:
+                    result = func(*args, **kwargs)
+                except Exception:
+                    drop_stack_frame()
+                    raise
             arg_map[RESULT_TOKEN] = result
 
             _check_specs(post_specs)

docs/index.rst

@@ -495,3 +495,19 @@ your type:
    :start-after: [custom_type]
    :end-before: [custom_type]
    :dedent:
+
+
+`check_shapes` in stack traces
+++++++++++++++++++++++++++++++
+
+If you use :mod:`check_shapes` consistently you will have a lot of functions wrapped in
+:func:`check_shapes`. This means that if you have an error many of the stack frames in the trace
+back would belong to :mod:`check_shapes`. For the sake of more readable error messages
+:mod:`check_shapes` has code to hide itself from trace backs. If you do not like this you can
+disable this behaviour with :func:`set_drop_frames`:
+
+
+.. literalinclude:: /examples/test_check_shapes_examples.py
+   :start-after: [disable_drop_frames]
+   :end-before: [disable_drop_frames]
+   :dedent: