solve by extension class of KeyError

Signed-off-by: dafnapension <[email protected]>

Author: dafnapension

src/unitxt/error_utils.py

@@ -109,7 +109,7 @@ def _get_existing_context(error: Exception):
             existing["context_object"],
             existing["context"],
         )
-    return str(error), None, {}
+    return error.original_error if type(error) == ExtKeyError else str(error), None, {}
 
 
 def _format_object_context(obj: Any) -> Optional[str]:
@@ -239,13 +239,22 @@ def _store_context_attributes(
         "original_message": original_message,
     }
     try:
-        error.original_error = type(error)(original_message)
+        error.original_error = (
+            original_message
+            if type(error) == KeyError
+            else type(error)(original_message)
+        )
     except (TypeError, ValueError):
         error.original_error = Exception(original_message)
     error.context_object = context_object
     error.context = context
 
 
+class ExtKeyError(KeyError):
+    def __str__(self):
+        return "\n" + self.args[0]
+
+
 def _add_context_to_exception(
     original_error: Exception, context_object: Any = None, **context
 ):
@@ -262,14 +271,17 @@ def _add_context_to_exception(
     }
     context_parts = _build_context_parts(final_context_object, final_context)
     context_message = _create_context_box(context_parts)
-    _store_context_attributes(
-        original_error, final_context_object, final_context, original_message
-    )
+    if type(original_error) == KeyError:
+        f = ExtKeyError(original_message)
+    else:
+        f = original_error
+    _store_context_attributes(f, final_context_object, final_context, original_message)
     if context_parts:
         formatted_message = f"\n{context_message}\n\n{original_message}"
-        original_error.args = (formatted_message,)
+        f.args = (formatted_message,)
     else:
-        original_error.args = (original_message,)
+        f.args = (original_message,)
+    return f
 
 
 @contextmanager
@@ -298,7 +310,5 @@ def error_context(context_object: Any = None, **context):
     try:
         yield
     except Exception as e:
-        if e.__class__.__name__ == "KeyError":
-            e = RuntimeError(e.__class__.__name__ + ": '" + e.args[0] + "'")
-        _add_context_to_exception(e, context_object, **context)
-        raise e
+        f = _add_context_to_exception(e, context_object, **context)
+        raise f from None

tests/library/test_error_utils.py

@@ -47,7 +47,7 @@ class TestProcessor:
 
         processor = TestProcessor()
 
-        with self.assertRaises(RuntimeError) as cm:
+        with self.assertRaises(KeyError) as cm:
             with error_context(processor):
                 raise KeyError("Missing key")
 
@@ -186,12 +186,12 @@ class TestProcessor:
 
     def test_error_context_without_object(self):
         """Test error_context without a context object."""
-        with self.assertRaises(RuntimeError) as cm:
+        with self.assertRaises(KeyError) as cm:
             with error_context(input_file="data.json", line_number=156):
                 raise KeyError("Missing field")
 
         error = cm.exception
-        self.assertIsInstance(error, RuntimeError)
+        self.assertIsInstance(error, KeyError)
         self.assertIsNone(error.context_object)
         # Context now includes version info plus the specified context
         self.assertIn("Unitxt", error.context)