image-match/python/py/Lib/site-packages/sentry_sdk/integrations/sys_exit.py

import functools
import sys

import sentry_sdk
from sentry_sdk.utils import capture_internal_exceptions, event_from_exception
from sentry_sdk.integrations import Integration
from sentry_sdk._types import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Callable
    from typing import NoReturn, Union


class SysExitIntegration(Integration):
    """Captures sys.exit calls and sends them as events to Sentry.

    By default, SystemExit exceptions are not captured by the SDK. Enabling this integration will capture SystemExit
    exceptions generated by sys.exit calls and send them to Sentry.

    This integration, in its default configuration, only captures the sys.exit call if the exit code is a non-zero and
    non-None value (unsuccessful exits). Pass `capture_successful_exits=True` to capture successful exits as well.
    Note that the integration does not capture SystemExit exceptions raised outside a call to sys.exit.
    """

    identifier = "sys_exit"

    def __init__(self, *, capture_successful_exits: bool = False) -> None:
        self._capture_successful_exits = capture_successful_exits

    @staticmethod
    def setup_once() -> None:
        SysExitIntegration._patch_sys_exit()

    @staticmethod
    def _patch_sys_exit() -> None:
        old_exit: "Callable[[Union[str, int, None]], NoReturn]" = sys.exit

        @functools.wraps(old_exit)
        def sentry_patched_exit(__status: "Union[str, int, None]" = 0) -> "NoReturn":
            # @ensure_integration_enabled ensures that this is non-None
            integration = sentry_sdk.get_client().get_integration(SysExitIntegration)
            if integration is None:
                old_exit(__status)

            try:
                old_exit(__status)
            except SystemExit as e:
                with capture_internal_exceptions():
                    if integration._capture_successful_exits or __status not in (
                        0,
                        None,
                    ):
                        _capture_exception(e)
                raise e

        sys.exit = sentry_patched_exit


def _capture_exception(exc: "SystemExit") -> None:
    event, hint = event_from_exception(
        exc,
        client_options=sentry_sdk.get_client().options,
        mechanism={"type": SysExitIntegration.identifier, "handled": False},
    )
    sentry_sdk.capture_event(event, hint=hint)