ref: ensure_integration_enabled without original function (getsentry#2893)

szokeasaurusrex · sentrivana · web-flow · commit 4fc60aaa5c7f · 2024-03-25T15:06:32.000+01:00
ensure_integration_enabled and ensure_integration_enabled_async can now decorate functions that return None without an original function.

---------

Co-authored-by: Ivana Kellyerova &lt;ivana.kellyerova@sentry.io&gt;
diff --git a/sentry_sdk/utils.py b/sentry_sdk/utils.py
@@ -38,12 +38,14 @@
     from typing import (
         Any,
         Callable,
+        cast,
         ContextManager,
         Dict,
         Iterator,
         List,
         NoReturn,
         Optional,
+        overload,
         ParamSpec,
         Set,
         Tuple,
@@ -1631,9 +1633,39 @@ def reraise(tp, value, tb=None):
     raise value
 
 
+def _no_op(*_a, **_k):
+    # type: (*Any, **Any) -> None
+    """No-op function for ensure_integration_enabled."""
+    pass
+
+
+async def _no_op_async(*_a, **_k):
+    # type: (*Any, **Any) -> None
+    """No-op function for ensure_integration_enabled_async."""
+    pass
+
+
+if TYPE_CHECKING:
+
+    @overload
+    def ensure_integration_enabled(
+        integration,  # type: type[sentry_sdk.integrations.Integration]
+        original_function,  # type: Callable[P, R]
+    ):
+        # type: (...) -> Callable[[Callable[P, R]], Callable[P, R]]
+        ...
+
+    @overload
+    def ensure_integration_enabled(
+        integration,  # type: type[sentry_sdk.integrations.Integration]
+    ):
+        # type: (...) -> Callable[[Callable[P, None]], Callable[P, None]]
+        ...
+
+
 def ensure_integration_enabled(
     integration,  # type: type[sentry_sdk.integrations.Integration]
-    original_function,  # type: Callable[P, R]
+    original_function=_no_op,  # type: Union[Callable[P, R], Callable[P, None]]
 ):
     # type: (...) -> Callable[[Callable[P, R]], Callable[P, R]]
     """
@@ -1657,25 +1689,51 @@ def patch_my_function():
             return my_function()
     ```
     """
+    if TYPE_CHECKING:
+        # Type hint to ensure the default function has the right typing. The overloads
+        # ensure the default _no_op function is only used when R is None.
+        original_function = cast(Callable[P, R], original_function)
 
     def patcher(sentry_patched_function):
         # type: (Callable[P, R]) -> Callable[P, R]
-        @wraps(original_function)
         def runner(*args: "P.args", **kwargs: "P.kwargs"):
             # type: (...) -> R
             if sentry_sdk.get_client().get_integration(integration) is None:
                 return original_function(*args, **kwargs)
 
             return sentry_patched_function(*args, **kwargs)
 
-        return runner
+        if original_function is _no_op:
+            return wraps(sentry_patched_function)(runner)
+
+        return wraps(original_function)(runner)
 
     return patcher
 
 
-def ensure_integration_enabled_async(
+if TYPE_CHECKING:
+
+    # mypy has some trouble with the overloads, hence the ignore[no-overload-impl]
+    @overload  # type: ignore[no-overload-impl]
+    def ensure_integration_enabled_async(
+        integration,  # type: type[sentry_sdk.integrations.Integration]
+        original_function,  # type: Callable[P, Awaitable[R]]
+    ):
+        # type: (...) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]
+        ...
+
+    @overload
+    def ensure_integration_enabled_async(
+        integration,  # type: type[sentry_sdk.integrations.Integration]
+    ):
+        # type: (...) -> Callable[[Callable[P, Awaitable[None]]], Callable[P, Awaitable[None]]]
+        ...
+
+
+# The ignore[no-redef] also needed because mypy is struggling with these overloads.
+def ensure_integration_enabled_async(  # type: ignore[no-redef]
     integration,  # type: type[sentry_sdk.integrations.Integration]
-    original_function,  # type: Callable[P, Awaitable[R]]
+    original_function=_no_op_async,  # type: Union[Callable[P, Awaitable[R]], Callable[P, Awaitable[None]]]
 ):
     # type: (...) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]
     """
@@ -1684,17 +1742,24 @@ def ensure_integration_enabled_async(
     Please refer to the `ensure_integration_enabled` documentation for more information.
     """
 
+    if TYPE_CHECKING:
+        # Type hint to ensure the default function has the right typing. The overloads
+        # ensure the default _no_op function is only used when R is None.
+        original_function = cast(Callable[P, Awaitable[R]], original_function)
+
     def patcher(sentry_patched_function):
         # type: (Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]
-        @wraps(original_function)
         async def runner(*args: "P.args", **kwargs: "P.kwargs"):
             # type: (...) -> R
             if sentry_sdk.get_client().get_integration(integration) is None:
                 return await original_function(*args, **kwargs)
 
             return await sentry_patched_function(*args, **kwargs)
 
-        return runner
+        if original_function is _no_op_async:
+            return wraps(sentry_patched_function)(runner)
+
+        return wraps(original_function)(runner)
 
     return patcher
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -594,6 +594,7 @@ def function_to_patch():
     )
 
     assert patched_function() == "patched"
+    assert patched_function.__name__ == "original_function"
 
 
 def test_ensure_integration_enabled_integration_disabled(sentry_init):
@@ -611,6 +612,41 @@ def function_to_patch():
     )
 
     assert patched_function() == "original"
+    assert patched_function.__name__ == "original_function"
+
+
+def test_ensure_integration_enabled_no_original_function_enabled(sentry_init):
+    shared_variable = "original"
+
+    def function_to_patch():
+        nonlocal shared_variable
+        shared_variable = "patched"
+
+    sentry_init(integrations=[TestIntegration])
+
+    # Test the decorator by applying to function_to_patch
+    patched_function = ensure_integration_enabled(TestIntegration)(function_to_patch)
+    patched_function()
+
+    assert shared_variable == "patched"
+    assert patched_function.__name__ == "function_to_patch"
+
+
+def test_ensure_integration_enabled_no_original_function_disabled(sentry_init):
+    shared_variable = "original"
+
+    def function_to_patch():
+        nonlocal shared_variable
+        shared_variable = "patched"
+
+    sentry_init(integrations=[])
+
+    # Test the decorator by applying to function_to_patch
+    patched_function = ensure_integration_enabled(TestIntegration)(function_to_patch)
+    patched_function()
+
+    assert shared_variable == "original"
+    assert patched_function.__name__ == "function_to_patch"
 
 
 @pytest.mark.asyncio
@@ -630,6 +666,7 @@ async def function_to_patch():
     )(function_to_patch)
 
     assert await patched_function() == "patched"
+    assert patched_function.__name__ == "original_function"
 
 
 @pytest.mark.asyncio
@@ -649,3 +686,48 @@ async def function_to_patch():
     )(function_to_patch)
 
     assert await patched_function() == "original"
+    assert patched_function.__name__ == "original_function"
+
+
+@pytest.mark.asyncio
+async def test_ensure_integration_enabled_async_no_original_function_enabled(
+    sentry_init,
+):
+    shared_variable = "original"
+
+    async def function_to_patch():
+        nonlocal shared_variable
+        shared_variable = "patched"
+
+    sentry_init(integrations=[TestIntegration])
+
+    # Test the decorator by applying to function_to_patch
+    patched_function = ensure_integration_enabled_async(TestIntegration)(
+        function_to_patch
+    )
+    await patched_function()
+
+    assert shared_variable == "patched"
+    assert patched_function.__name__ == "function_to_patch"
+
+
+@pytest.mark.asyncio
+async def test_ensure_integration_enabled_async_no_original_function_disabled(
+    sentry_init,
+):
+    shared_variable = "original"
+
+    async def function_to_patch():
+        nonlocal shared_variable
+        shared_variable = "patched"
+
+    sentry_init(integrations=[])
+
+    # Test the decorator by applying to function_to_patch
+    patched_function = ensure_integration_enabled_async(TestIntegration)(
+        function_to_patch
+    )
+    await patched_function()
+
+    assert shared_variable == "original"
+    assert patched_function.__name__ == "function_to_patch"
diff --git a/tests/tracing/test_decorator.py b/tests/tracing/test_decorator.py
@@ -15,6 +15,7 @@ async def my_async_example_function():
     return "return_of_async_function"
 
 
+@pytest.mark.forked
 def test_trace_decorator():
     with patch_start_tracing_child() as fake_start_child:
         result = my_example_function()
