Use appropriate error (KeyError) in HandlerTypeRegistry.get().

PiperOrigin-RevId: 718148036
google · Jan 22, 2025 · f7bbe80 · f7bbe80
1 parent fb272d0
commit f7bbe80
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/checkpoint/orbax/checkpoint/_src/handlers/handler_type_registry.py b/checkpoint/orbax/checkpoint/_src/handlers/handler_type_registry.py
@@ -83,7 +83,7 @@ def get(
   ) -> Type[CheckpointHandler]:
     """Gets an entry from the registry."""
     if handler_typestr not in self._registry:
-      raise ValueError(
+      raise KeyError(
           f'Handler type string "{handler_typestr}" not found in the registry.'
       )
     return self._registry[handler_typestr]

diff --git a/checkpoint/orbax/checkpoint/_src/handlers/handler_type_registry_test.py b/checkpoint/orbax/checkpoint/_src/handlers/handler_type_registry_test.py
@@ -158,7 +158,7 @@ def test_register_duplicate_handler_type(self):
   def test_get_handler_type_not_found(self):
     registry = HandlerTypeRegistry()
     with self.assertRaisesRegex(
-        ValueError,
+        KeyError,
         'Handler type string '
         r'"(?:__main__|handler_type_registry_test)\.TestHandler"'
         ' not found in the registry.',