Add typeguard in open-source

PiperOrigin-RevId: 726396947

Author: Conchylicultor

.github/workflows/pytest_and_autopublish.yml

@@ -39,7 +39,6 @@ jobs:
     # * sweep_utils_test: Depends on kxm
     # * lpips_test: Missing VGG weights
     # * partial_loader_test: Orbax partial checkpoint loader not yet open-sourced  (TODO(epot): Restore)
-    # * typing tests: Not yet supported due to typeguard version issues.
     - name: Run core tests
       run: |
         pytest -vv -n auto \
@@ -48,9 +47,7 @@ jobs:
           --ignore=kauldron/xm/ \
           --ignore=kauldron/metrics/lpips_test.py \
           --ignore=kauldron/checkpoints/partial_loader_test.py \
-          --ignore=kauldron/utils/sweep_utils_test.py \
-          --ignore=kauldron/typing/shape_spec_test.py \
-          --ignore=kauldron/typing/type_check_test.py
+          --ignore=kauldron/utils/sweep_utils_test.py
 
   # Auto-publish when version is increased
   publish-job:

kauldron/typing/shape_spec.py

@@ -83,8 +83,8 @@ def _assert_caller_is_typechecked_func() -> None:
   if stack[i + 1].function != "_reraise_with_shape_info":
     caller_name = stack[i].function
     raise AssertionError(
-        "Dim and Shape not yet supported due to `typeguard` issue."
-        f" Raised in {caller_name!r}"
+        "Dim and Shape only work inside of @typechecked functions. But"
+        f" {caller_name!r} lacks @typechecked."
     )
 
 

kauldron/typing/type_check.py

@@ -46,13 +46,10 @@ def check_type(
     expected_type: Any,
 ) -> None:
   """Ensure that value matches expected_type, alias for typeguard.check_type."""
-  if True:  # Typeguard not yet supported
-    return
   return typeguard.check_type(value, expected_type)
 
 
-exc_cls = Exception
-class TypeCheckError(exc_cls):
+class TypeCheckError(typeguard.TypeCheckError):
   """Indicates a runtime typechecking error from the @typechecked decorator."""
 
   def __init__(
@@ -112,9 +109,6 @@ def _annotation_repr(ann: Any) -> str:
 
 def typechecked(fn):
   """Decorator to enable runtime type-checking and shape-checking."""
-  if True:  # Typeguard not yet supported
-    return fn
-
   if hasattr(fn, "__wrapped__"):
     raise AssertionError("@typechecked should be the innermost decorator")
 
@@ -469,3 +463,7 @@ def add_custom_checker_lookup_fn(lookup_fn):
       break
   else:  # prepend
     checker_lookup_fns[:0] = [lookup_fn]
+
+
+add_custom_checker_lookup_fn(_array_spec_checker_lookup)
+add_custom_checker_lookup_fn(_dataclass_checker_lookup)

kauldron/typing/type_check_test.py

@@ -16,7 +16,7 @@
 
 import jaxtyping as jt
 from kauldron.typing import Float, TypeCheckError, typechecked  # pylint: disable=g-multiple-import,g-importing-member
-from kauldron.typing import type_check
+from kauldron.typing import type_check  # pylint: disable=g-bad-import-order
 import numpy as np
 import pytest
 

pyproject.toml

@@ -44,7 +44,7 @@ dependencies = [
     "tqdm",
     # TODO(klausg): Restore typeguard or switch to something else
     # closest match to the internal typeguard
-    # "typeguard@git+https://github.com/agronholm/typeguard@0dd7f7510b7c694e66a0d17d1d58d185125bad5d",
+    "typeguard",
     "typing_extensions",
     "xmanager",
     # lazy deps (should ideally remove those)