diff --git a/dmf/alerts/alerts.py b/dmf/alerts/alerts.py index fea1aec..d7e30ee 100644 --- a/dmf/alerts/alerts.py +++ b/dmf/alerts/alerts.py @@ -1,5 +1,5 @@ from typing import Optional, TYPE_CHECKING, Tuple, Union - +import warnings from ..utils.typing import Literal from ..env import env @@ -58,6 +58,11 @@ def send_message( """ backend = get_backend() + + if not backend: + warnings.warn("No alert backend is available. Message not sent.") + return + backend.send_message( text=text, attachment=attachment, scheduled_time=scheduled_time ) @@ -102,17 +107,20 @@ def send_alert( Sending an error alert with an attachment:: - send_alert(text="An error occurred. See the attached log file.", attachment="/path/to/log.txt", level="error") + send_alert(text="An error occurred. See the attached log file.", attachmenlogging.warning("No alert backend is available. Message not sent.")t="/path/to/log.txt", level="error") """ # Get or create the backend instance (always storing it globally) backend = get_backend() + if not backend: + warnings.warn("No alert backend is available. Message not sent.") + return backend.send_alert(text=text, attachment=attachment, params=params, level=level) def get_backend( alert_token: Optional[str] = None, store: bool = True -) -> "AlertBackend": +) -> Optional["AlertBackend"]: """ Get or create an AlertBackend instance. If the global backend is None or store is False, initialize it with resolved credentials. @@ -135,7 +143,7 @@ def get_backend( current_backend = TelegramBackend(token=alert_token) else: - raise ValueError(f"Unsupported credential type: {credential_type}") + current_backend = None # Store the backend if required if store: @@ -149,7 +157,7 @@ def get_backend( def resolve_credentials( alert_token: Optional[str] = None, -) -> Tuple[str, Literal["slack", "telegram"]]: +) -> Tuple[Optional[str], Optional[Literal["slack", "telegram"]]]: """ Resolve the credentials for the alert system. Look for environment variables, or use provided overrides. @@ -159,9 +167,7 @@ def resolve_credentials( alert_token = alert_token or env.getenv(ALERT_TOKEN) if not alert_token: - raise ValueError( - f"Alert token must be provided or set in the {ALERT_TOKEN} environment variable." - ) + return None, None if alert_token.startswith("xoxb-"): credential_type = "slack" @@ -169,5 +175,7 @@ def resolve_credentials( elif len(alert_token.split(":")) == 2: credential_type = "telegram" + else: + credential_type = None return alert_token, credential_type diff --git a/dmf/alerts/decorator.py b/dmf/alerts/decorator.py index a04351c..54e7dda 100644 --- a/dmf/alerts/decorator.py +++ b/dmf/alerts/decorator.py @@ -2,9 +2,14 @@ from datetime import datetime from pathlib import Path + from .alerts import send_alert -__all__ = ["alert"] +__all__ = ["alert", "DoNotAlert"] + +class DoNotAlert(Exception): + """Exception raised to prevent an alert from being sent.""" + pass def alert( @@ -13,6 +18,7 @@ def alert( output: bool = False, input: bool = False, max_length: int = 100, + disable: bool = False, ) -> Callable: """ Decorator that sends an alert after the execution of the decorated function. @@ -64,6 +70,12 @@ def hello(name): print(f"Hello {name}") return Path("/path/to/{name}.pdf") """ + # Quick disable of the alert decorator + if disable: + if func is None: + return lambda func: func + return func + if func is None: # Used as @alert(output=True, input=True) with parentheses def decorator(inner_func: Callable) -> Callable: @@ -76,7 +88,6 @@ def decorator(inner_func: Callable) -> Callable: # Used as @alert without parentheses return _alert_decorator(func, output=output, input=input, max_length=max_length) - def _alert_decorator( func: Callable, output: bool, @@ -118,6 +129,9 @@ def wrapper(*args, **kwargs): text=message, level="success", params=params, attachment=attachment ) return result + except DoNotAlert: + # Do not send an alert if the DoNotAlert exception is raised + return result except Exception as e: end_time = datetime.now() message = f"Execution of function *{func.__name__}* failed with an error." @@ -148,3 +162,6 @@ def _input_to_str(args, kwargs, input, max_length: int = 100) -> str: if input is True or key in input: input_str += f"{sep}*{key}*: _{_to_str(value, max_length)}_" return input_str + +# Assign the DoNotAlert exception to the alert decorator +setattr(alert, "DoNotAlert", DoNotAlert) \ No newline at end of file