1- import sys
2- from contextlib import contextmanager
1+ from functools import wraps
2+ from inspect import iscoroutinefunction
33
44from sentry_sdk ._types import TYPE_CHECKING
55from sentry_sdk .crons import capture_checkin
66from sentry_sdk .crons .consts import MonitorStatus
7- from sentry_sdk .utils import now , reraise
7+ from sentry_sdk .utils import now
88
99if TYPE_CHECKING :
10- from typing import Generator , Optional
10+ from types import TracebackType
11+ from typing import (
12+ Awaitable ,
13+ Callable ,
14+ Optional ,
15+ ParamSpec ,
16+ Type ,
17+ TypeVar ,
18+ Union ,
19+ )
20+
21+ P = ParamSpec ("P" )
22+ R = TypeVar ("R" )
1123
1224
13- @contextmanager
14- def monitor (monitor_slug = None ):
15- # type: (Optional[str]) -> Generator[None, None, None]
25+ class monitor : # noqa: N801
1626 """
1727 Decorator/context manager to capture checkin events for a monitor.
1828
@@ -39,32 +49,51 @@ def test(arg):
3949 with sentry_sdk.monitor(monitor_slug='my-fancy-slug'):
4050 print(arg)
4151 ```
52+ """
4253
54+ def __init__ (self , monitor_slug = None ):
55+ # type: (Optional[str]) -> None
56+ self .monitor_slug = monitor_slug
4357
44- """
58+ def __enter__ (self ):
59+ # type: () -> None
60+ self .start_timestamp = now ()
61+ self .check_in_id = capture_checkin (
62+ monitor_slug = self .monitor_slug , status = MonitorStatus .IN_PROGRESS
63+ )
4564
46- start_timestamp = now ()
47- check_in_id = capture_checkin (
48- monitor_slug = monitor_slug , status = MonitorStatus .IN_PROGRESS
49- )
65+ def __exit__ (self , exc_type , exc_value , traceback ):
66+ # type: (Optional[Type[BaseException]], Optional[BaseException], Optional[TracebackType]) -> None
67+ duration_s = now () - self .start_timestamp
68+
69+ if exc_type is None and exc_value is None and traceback is None :
70+ status = MonitorStatus .OK
71+ else :
72+ status = MonitorStatus .ERROR
5073
51- try :
52- yield
53- except Exception :
54- duration_s = now () - start_timestamp
5574 capture_checkin (
56- monitor_slug = monitor_slug ,
57- check_in_id = check_in_id ,
58- status = MonitorStatus . ERROR ,
75+ monitor_slug = self . monitor_slug ,
76+ check_in_id = self . check_in_id ,
77+ status = status ,
5978 duration = duration_s ,
6079 )
61- exc_info = sys .exc_info ()
62- reraise (* exc_info )
63-
64- duration_s = now () - start_timestamp
65- capture_checkin (
66- monitor_slug = monitor_slug ,
67- check_in_id = check_in_id ,
68- status = MonitorStatus .OK ,
69- duration = duration_s ,
70- )
80+
81+ def __call__ (self , fn ):
82+ # type: (Callable[P, R]) -> Callable[P, Union[R, Awaitable[R]]]
83+ if iscoroutinefunction (fn ):
84+
85+ @wraps (fn )
86+ async def inner (* args : "P.args" , ** kwargs : "P.kwargs" ):
87+ # type: (...) -> R
88+ with self :
89+ return await fn (* args , ** kwargs )
90+
91+ else :
92+
93+ @wraps (fn )
94+ def inner (* args : "P.args" , ** kwargs : "P.kwargs" ):
95+ # type: (...) -> R
96+ with self :
97+ return fn (* args , ** kwargs )
98+
99+ return inner
0 commit comments