Fix UNKNOWN_DIM formatting in ktyping error messages.

PiperOrigin-RevId: 889752042

Author: Qwlouse

kauldron/ktyping/dim_view.py

@@ -22,6 +22,7 @@
 from kauldron.ktyping import errors
 from kauldron.ktyping import internal_typing
 from kauldron.ktyping import scope as kscope
+from kauldron.ktyping import utils
 
 DimValue = internal_typing.DimValue
 DimValues = internal_typing.DimValues
@@ -217,19 +218,7 @@ def __str__(self) -> str:
 
 def _format_dim_value(value: DimValue) -> str:
   """Formats a DimValue tuple into a human-readable string."""
-
-  def _fmt(v):
-    if isinstance(v, int):
-      return str(v)
-    if v == UNKNOWN_DIM:
-      return "#"
-    return f"&{v}"
-
-  str_values = [_fmt(v) for v in value]
-  if len(value) == 1:
-    return str_values[0]
-  else:
-    return f"({', '.join(str_values)})"
+  return utils.format_dim_value(value)
 
 
 def _format_dim_assignment(dim_name: str, value: DimValue, align: int) -> str:

kauldron/ktyping/errors.py

@@ -204,9 +204,9 @@ def __str__(self):
 
 def _format_dim_assignment(dim, value):
   if len(value) == 1:
-    return f"{dim}: {value[0]}"
+    return f"{dim}: {utils.format_dim_value(value)}"
   else:
-    return f"*{dim}: {value}"
+    return f"*{dim}: {utils.format_dim_value(value)}"
 
 
 # MARK: Error messages

kauldron/ktyping/errors_test.py

@@ -61,3 +61,55 @@ def test_array_type_check():
 
   assert exc.candidates_block == "Dim Assignments:\n - {a: 2, b: 3}"
   assert exc.candidates_block in msg
+
+
+def test_format_dim_assignment_with_unknown_dim():
+  """UNKNOWN_DIM should be formatted as '#' in dim assignments."""
+  from kauldron.ktyping import internal_typing  # pylint: disable=g-import-not-at-top
+
+  unknown = internal_typing.UNKNOWN_DIM
+
+  # Single unknown dim
+  result = errors._format_dim_assignment("batch", (unknown,))
+  assert result == "batch: #", result
+
+  # Multi-dim with unknown
+  result = errors._format_dim_assignment("data", (6, unknown))
+  assert result == "*data: (6, #)", result
+
+  # Multi-dim all unknown
+  result = errors._format_dim_assignment("b", (unknown, unknown))
+  assert result == "*b: (#, #)", result
+
+  # No unknown (sanity check)
+  result = errors._format_dim_assignment("x", (3, 5))
+  assert result == "*x: (3, 5)", result
+
+
+@typechecked
+def _broadcastable_fn(
+    x: Float["*#batch n"],
+    y: Int["*#batch m"],
+) -> Float["n"]:
+  del y
+  return x[..., :, 0]  # wrong return shape to trigger error
+
+
+def test_unknown_dim_in_error_message():
+  """UNKNOWN_DIM in candidates_block should show '#' not the enum repr."""
+  x = np.zeros((1, 6, 4), dtype=np.float32)  # #*batch=(1, 6) -> (#, 6)
+  y = np.zeros((1, 6, 3), dtype=np.int32)
+
+  with pytest.raises(errors.KTypeCheckError) as exc_info:
+    _broadcastable_fn(x, y)
+
+  msg = str(exc_info.value)
+  # The ugly enum repr should NOT appear in the error message
+  assert (
+      "UnknownDim" not in msg
+  ), f"UnknownDim leaked into error message:\n{msg}"
+  assert (
+      "UNKNOWN_DIM" not in msg
+  ), f"UNKNOWN_DIM leaked into error message:\n{msg}"
+  # The '#' formatting should be used instead
+  assert "#" in exc_info.value.candidates_block

kauldron/ktyping/utils.py

@@ -111,6 +111,24 @@ def _get_array_type_shorthand(array: Any) -> str:
   return get_type_name(array)
 
 
+# MARK: format_dim_value
+def format_dim_value(value: internal_typing.DimValue) -> str:
+  """Formats a DimValue tuple into a human-readable string."""
+
+  def _fmt(v):
+    if isinstance(v, int):
+      return str(v)
+    if v == internal_typing.UNKNOWN_DIM:
+      return "#"
+    return f"&{v}"
+
+  str_values = [_fmt(v) for v in value]
+  if len(value) == 1:
+    return str_values[0]
+  else:
+    return f"({', '.join(str_values)})"
+
+
 # MARK: get_type_hints
 def get_type_hints(fn: Callable[..., Any]) -> dict[str, Any]:
   """Return the type hints for the given function with caching."""