py-healthchecks.io/src/healthchecks_io/client/check_trap.py

"""CheckTrap is a context manager to wrap around python code to communicate results to a Healthchecks check."""

from types import TracebackType
from typing import List
from typing import Optional
from typing import Type
from typing import Union

from .async_client import AsyncClient
from .exceptions import PingFailedError
from .exceptions import WrongClientError
from .sync_client import Client


class CheckTrap:
    """CheckTrap is a context manager to wrap around python code to communicate results to a Healthchecks check."""

    def __init__(
        self,
        client: Union[Client, AsyncClient],
        uuid: str = "",
        slug: str = "",
        suppress_exceptions: bool = False,
    ) -> None:
        """A context manager to wrap around python code to communicate results to a Healthchecks check.

        Args:
            client (Union[Client, AsyncClient]): healthchecks_io client, async or sync
            uuid (str): uuid of the check. Defaults to "".
            slug (str): slug of the check, exclusion wiht uuid. Defaults to "".
            suppress_exceptions (bool): If true, do not raise any exceptions. Defaults to False.

        Raises:
            Exception: Raised if a slug and a uuid is passed
        """
        if uuid == "" and slug == "":
            raise Exception("Must pass a slug or an uuid")
        self.client: Union[Client, AsyncClient] = client
        self.uuid: str = uuid
        self.slug: str = slug
        self.log_lines: List[str] = list()
        self.suppress_exceptions: bool = suppress_exceptions

    def add_log(self, line: str) -> None:
        """Add a line to the context manager's log that is sent with the check.

        Args:
            line (str): String to add to the logs
        """
        self.log_lines.append(line)

    def __enter__(self) -> "CheckTrap":
        """Enter the context manager.

        Sends a start ping to the check represented by self.uuid or self.slug.

        Raises:
            WrongClientError: Raised when using an AsyncClient with this as a sync client manager
            PingFailedError: When a ping fails for any reason not handled by a custom exception
            HCAPIAuthError: Raised when status_code == 401 or 403
            HCAPIError: Raised when status_code is 5xx
            CheckNotFoundError: Raised when status_code is 404 or response text has "not found" in it
            BadAPIRequestError: Raised when status_code is 400, or if you pass a uuid and a slug, or if
                pinging by a slug and do not have a ping key set
            HCAPIRateLimitError: Raised when status code is 429 or response text has "rate limited" in it
            NonUniqueSlugError: Raused when status code is 409.

        Returns:
            CheckTrap: self
        """
        if isinstance(self.client, AsyncClient):
            raise WrongClientError("You passed an AsyncClient, use this as an async context manager")
        result = self.client.start_ping(uuid=self.uuid, slug=self.slug)
        if not result[0]:
            raise PingFailedError(result[1])
        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> Optional[bool]:
        """Exit the context manager.

        If there is an exception, add it to any log lines and send a fail ping.
        Otherwise, send a success ping with any log lines appended.

        Args:
            exc_type (Optional[Type[BaseException]]): [description]
            exc (Optional[BaseException]): [description]
            traceback (Optional[TracebackType]): [description]

        Returns:
            Optional[bool]: self.suppress_exceptions, if true will not raise any exceptions
        """
        if exc_type is None:
            self.client.success_ping(self.uuid, self.slug, data="\n".join(self.log_lines))
        else:
            self.add_log(str(exc))
            self.add_log(str(traceback))
            self.client.fail_ping(self.uuid, self.slug, data="\n".join(self.log_lines))
        return self.suppress_exceptions

    async def __aenter__(self) -> "CheckTrap":
        """Enter the context manager.

        Sends a start ping to the check represented by self.uuid or self.slug.

        Raises:
            WrongClientError: Raised when using an AsyncClient with this as a sync client manager
            PingFailedError: When a ping fails for any reason not handled by a custom exception
            HCAPIAuthError: Raised when status_code == 401 or 403
            HCAPIError: Raised when status_code is 5xx
            CheckNotFoundError: Raised when status_code is 404 or response text has "not found" in it
            BadAPIRequestError: Raised when status_code is 400, or if you pass a uuid and a slug, or if
                pinging by a slug and do not have a ping key set
            HCAPIRateLimitError: Raised when status code is 429 or response text has "rate limited" in it
            NonUniqueSlugError: Raused when status code is 409.

        Returns:
            CheckTrap: self
        """
        if isinstance(self.client, Client):
            raise WrongClientError("You passed a sync Client, use this as a regular context manager")
        result = await self.client.start_ping(self.uuid, self.slug)
        if not result[0]:
            raise PingFailedError(result[1])
        return self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> Optional[bool]:
        """Exit the context manager.

        If there is an exception, add it to any log lines and send a fail ping.
        Otherwise, send a success ping with any log lines appended.

        Args:
            exc_type (Optional[Type[BaseException]]): [description]
            exc (Optional[BaseException]): [description]
            traceback (Optional[TracebackType]): [description]

        Returns:
            Optional[bool]: self.suppress_exceptions, if true will not raise any exceptions
        """
        if exc_type is None:
            # ignore typing, if we've gotten here we know its an async client
            await self.client.success_ping(  # type: ignore
                self.uuid, self.slug, data="\n".join(self.log_lines)
            )
        else:
            self.add_log(str(exc))
            self.add_log(str(traceback))
            await self.client.fail_ping(  # type: ignore
                self.uuid, self.slug, data="\n".join(self.log_lines)
            )
        return self.suppress_exceptions