From 0a76c9c70c518f2ea2de4e95a87634d557bf7ba1 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Thu, 2 Nov 2023 20:01:27 -0700
Subject: [PATCH] Error message improvements

---
 jaxtyping/_decorator.py | 18 +++++++++++++-----
 jaxtyping/_raise.py     |  1 +
 test/test_decorator.py  |  2 +-
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py
index d16b1ab..0bfa5c8 100644
--- a/jaxtyping/_decorator.py
+++ b/jaxtyping/_decorator.py
@@ -687,17 +687,25 @@ def _exc_shape_info(memos) -> str:
     Used in type-checking error messages.
     """
     single_memo, variadic_memo, pytree_memo, _ = memos
+    single_memo = {
+        name: size
+        for name, size in single_memo.items()
+        if not name.startswith("~~delete~~")
+    }
+    variadic_memo = {
+        name: shape
+        for name, (_, shape) in variadic_memo.items()
+        if not name.startswith("~~delete~~")
+    }
     pieces = []
     if len(single_memo) > 0 or len(variadic_memo) > 0:
         pieces.append(
             "The current values for each jaxtyping axis annotation are as follows."
         )
         for name, size in single_memo.items():
-            if not name.startswith("~~delete~~"):
-                pieces.append(f"{name}={size}")
-        for name, (_, shape) in variadic_memo.items():
-            if not name.startswith("~~delete~~"):
-                pieces.append(f"{name}={shape}")
+            pieces.append(f"{name}={size}")
+        for name, shape in variadic_memo.items():
+            pieces.append(f"{name}={shape}")
     if len(pytree_memo) > 0:
         pieces.append(
             "The current values for each jaxtyping PyTree structure annotation are as "
diff --git a/jaxtyping/_raise.py b/jaxtyping/_raise.py
index 9230877..38d3f1a 100644
--- a/jaxtyping/_raise.py
+++ b/jaxtyping/_raise.py
@@ -6,6 +6,7 @@ def jaxtyping_raise(e) -> NoReturn:
     `TypeCheckError`. All `raise` statements from within `__instancecheck__` should use
     this.
     """
+    __tracebackhide__ = True
     try:
         raise e
     except Exception as f:
diff --git a/test/test_decorator.py b/test/test_decorator.py
index 3d0001f..52c0c9a 100644
--- a/test/test_decorator.py
+++ b/test/test_decorator.py
@@ -110,7 +110,7 @@ def test_varkwargs(jaxtyp, typecheck):
 
 def test_defaults(jaxtyp, typecheck):
     @jaxtyp(typecheck)
-    def f(x, y=1):
+    def f(x: int, y=1):
         pass
 
     f(1)