18
18
from typing_extensions import Protocol , TypeVar , runtime_checkable
19
19
20
20
import taskiq .exceptions # noqa: WPS301
21
+ from taskiq .compat import IS_PYDANTIC2 , validate_call
21
22
22
23
DecodedType = TypeVar ("DecodedType" )
23
24
EncodedType = TypeVar ("EncodedType" )
@@ -245,18 +246,34 @@ def get_pickled_exception(exc: BaseException) -> BaseException:
245
246
return exc
246
247
247
248
248
- class ExceptionRepr (pydantic .BaseModel ):
249
- """Serialiable exception representation."""
249
+ if IS_PYDANTIC2 :
250
250
251
- exc_type : str
252
- exc_message : Tuple [Any , ...]
253
- exc_module : Optional [str ]
254
- exc_cause : Optional [Union [BaseException , "ExceptionRepr" ]] = None
255
- exc_context : Optional [Union [BaseException , "ExceptionRepr" ]] = None
256
- exc_suppress_context : bool = False
251
+ class ExceptionRepr (pydantic .BaseModel ):
252
+ """Serializable exception model for pydantic v2."""
257
253
258
- class Config :
259
- arbitrary_types_allowed = True
254
+ exc_type : str
255
+ exc_message : Tuple [Any , ...]
256
+ exc_module : Optional [str ]
257
+ exc_cause : Optional [Union [BaseException , "ExceptionRepr" ]] = None
258
+ exc_context : Optional [Union [BaseException , "ExceptionRepr" ]] = None
259
+ exc_suppress_context : bool = False
260
+
261
+ model_config = pydantic .ConfigDict (arbitrary_types_allowed = True )
262
+
263
+ else :
264
+
265
+ class ExceptionRepr (pydantic .BaseModel ): # type: ignore
266
+ """Serializable exception model for pydantic v1."""
267
+
268
+ exc_type : str
269
+ exc_message : Tuple [Any , ...]
270
+ exc_module : Optional [str ]
271
+ exc_cause : Optional [Union [BaseException , "ExceptionRepr" ]] = None
272
+ exc_context : Optional [Union [BaseException , "ExceptionRepr" ]] = None
273
+ exc_suppress_context : bool = False
274
+
275
+ class Config :
276
+ arbitrary_types_allowed = True
260
277
261
278
262
279
def _prepare_exception (
@@ -297,7 +314,7 @@ def _prepare_exception(
297
314
SEEN_EXCEPTIONS_CACHE .discard (id (exc ))
298
315
299
316
300
- @pydantic . validate_arguments (config = { " arbitrary_types_allowed" : True }) # type: ignore
317
+ @validate_call (config = pydantic . ConfigDict ( arbitrary_types_allowed = True ))
301
318
def prepare_exception (
302
319
exc : BaseException ,
303
320
coder : Coder [Any , Any ],
@@ -312,7 +329,7 @@ def prepare_exception(
312
329
return _prepare_exception (exc , coder ) # type: ignore
313
330
314
331
315
- @pydantic . validate_arguments (config = { " arbitrary_types_allowed" : True }) # type: ignore
332
+ @validate_call (config = pydantic . ConfigDict ( arbitrary_types_allowed = True ))
316
333
def exception_to_python ( # noqa: C901, WPS210
317
334
exc : Optional [Union [BaseException , ExceptionRepr ]],
318
335
) -> Optional [BaseException ]:
0 commit comments