Coverage for src/zapy/requests/exceptions.py: 100%

36 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-02-10 19:35 +0000

1import asyncio 

2import functools 

3from typing import Any, Callable, TypedDict, TypeVar, cast 

4 

5from zapy.base.exceptions import HandledError, ZapyError 

6from zapy.templating.traceback import TracebackInfo, copy_traceback 

7 

8from .models import HttpxResponse 

9 

10 

11class RenderLocationContext(TypedDict): 

12 stacktrace: TracebackInfo 

13 response: HttpxResponse 

14 

15 

16class RenderLocationError(ZapyError, HandledError): 

17 

18 namespace = "render" 

19 context: RenderLocationContext | dict[str, Any] # type: ignore[assignment] 

20 

21 def __init__(self, ex: Exception, location: str): 

22 super().__init__(f"Error on {location}") 

23 self.identifier = location 

24 info = copy_traceback(self, from_exc=ex) 

25 if info: 

26 self.context["stacktrace"] = info 

27 

28 

29T = TypeVar("T", bound=Callable[..., Any]) 

30 

31 

32def error_location(location: str) -> Callable[[T], T]: 

33 def decorator(function: T) -> T: 

34 if not asyncio.iscoroutinefunction(function): 

35 

36 @functools.wraps(function) 

37 def wrapper(*args: Any, **kwargs: Any) -> Any: 

38 try: 

39 return function(*args, **kwargs) 

40 except Exception as ex: 

41 raise RenderLocationError(ex, location) from ex 

42 

43 else: 

44 

45 @functools.wraps(function) 

46 async def wrapper(*args: Any, **kwargs: Any) -> Any: 

47 try: 

48 return await function(*args, **kwargs) 

49 except Exception as ex: 

50 raise RenderLocationError(ex, location) from ex 

51 

52 return cast(T, wrapper) 

53 

54 return decorator