Skip to content

Commit adc56bd

Browse files
authored
chore(internal): add is_wrapped and is_wrapped_with helpers (#14295)
Useful helpers to ensure we can check if something is already wrapped to avoid re-wrapping it, and we can check if it is already wrapped with a specific function. Useful both for testing and patching. ## Checklist - [x] PR author has checked that all the criteria below are met - The PR description includes an overview of the change - The PR description articulates the motivation for the change - The change includes tests OR the PR description describes a testing strategy - The PR description notes risks associated with the change, if any - Newly-added code is easy to change - The change follows the [library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) - The change includes or references documentation updates if necessary - Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) ## Reviewer Checklist - [ ] Reviewer has checked that all the criteria below are met - Title is accurate - All changes are related to the pull request's stated goal - Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - Testing strategy adequately addresses listed risks - Newly-added code is easy to change - Release note makes sense to a user of the library - If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
1 parent 1d248bd commit adc56bd

File tree

2 files changed

+99
-0
lines changed

2 files changed

+99
-0
lines changed

ddtrace/internal/wrapping/__init__.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,38 @@ def wrap(f, wrapper):
269269
return wf
270270

271271

272+
def is_wrapped(f: FunctionType) -> bool:
273+
"""Check if a function is wrapped with any wrapper."""
274+
try:
275+
wf = cast(WrappedFunction, f)
276+
inner = cast(FunctionType, wf.__dd_wrapped__)
277+
278+
# Sanity check
279+
assert inner.__name__ == "<wrapped>", "Wrapper has wrapped function" # nosec
280+
return True
281+
except AttributeError:
282+
return False
283+
284+
285+
def is_wrapped_with(f: FunctionType, wrapper: Wrapper) -> bool:
286+
"""Check if a function is wrapped with a specific wrapper."""
287+
try:
288+
wf = cast(WrappedFunction, f)
289+
inner = cast(FunctionType, wf.__dd_wrapped__)
290+
291+
# Sanity check
292+
assert inner.__name__ == "<wrapped>", "Wrapper has wrapped function" # nosec
293+
294+
if wrapper in f.__code__.co_consts:
295+
return True
296+
297+
# This is not the correct wrapping layer. Try with the next one.
298+
return is_wrapped_with(inner, wrapper)
299+
300+
except AttributeError:
301+
return False
302+
303+
272304
def unwrap(wf, wrapper):
273305
# type: (WrappedFunction, Wrapper) -> FunctionType
274306
"""Unwrap a wrapped function.

tests/internal/test_wrapping.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66

77
import pytest
88

9+
from ddtrace.internal.wrapping import is_wrapped
10+
from ddtrace.internal.wrapping import is_wrapped_with
911
from ddtrace.internal.wrapping import unwrap
1012
from ddtrace.internal.wrapping import wrap
1113
from ddtrace.internal.wrapping.context import WrappingContext
@@ -95,6 +97,71 @@ def f(a, b, c=None):
9597
assert not channel1 and not channel2
9698

9799

100+
def test_is_wrapped():
101+
"""Test that `is_wrapped` and `is_wrapped_with` work as expected."""
102+
103+
def first_wrapper(f, args, kwargs):
104+
return f(*args, **kwargs)
105+
106+
def second_wrapper(f, args, kwargs):
107+
return f(*args, **kwargs)
108+
109+
def f(a, b, c=None):
110+
return (a, b, c)
111+
112+
# Function works
113+
assert f(1, 2) == (1, 2, None)
114+
115+
# Not wrapped yet
116+
assert not is_wrapped(f)
117+
assert not is_wrapped_with(f, first_wrapper)
118+
assert not is_wrapped_with(f, second_wrapper)
119+
120+
# Wrap with first wrapper
121+
wrap(f, first_wrapper)
122+
123+
# Function still works
124+
assert f(1, 2) == (1, 2, None)
125+
126+
# Only wrapped with first_wrapper
127+
assert is_wrapped(f)
128+
assert is_wrapped_with(f, first_wrapper)
129+
assert not is_wrapped_with(f, second_wrapper)
130+
131+
# Wrap with second wrapper
132+
wrap(f, second_wrapper)
133+
134+
# Function still works
135+
assert f(1, 2) == (1, 2, None)
136+
137+
# Wrapped with everything
138+
assert is_wrapped(f)
139+
assert is_wrapped_with(f, first_wrapper)
140+
assert is_wrapped_with(f, second_wrapper)
141+
142+
# Unwrap first wrapper
143+
unwrap(f, first_wrapper)
144+
145+
# Function still works
146+
assert f(1, 2) == (1, 2, None)
147+
148+
# Still wrapped with second_wrapper
149+
assert is_wrapped(f)
150+
assert not is_wrapped_with(f, first_wrapper)
151+
assert is_wrapped_with(f, second_wrapper)
152+
153+
# Unwrap second wrapper
154+
unwrap(f, second_wrapper)
155+
156+
# Function still works
157+
assert f(1, 2) == (1, 2, None)
158+
159+
# Not wrapped anymore
160+
assert not is_wrapped(f)
161+
assert not is_wrapped_with(f, first_wrapper)
162+
assert not is_wrapped_with(f, second_wrapper)
163+
164+
98165
@pytest.mark.skipif(sys.version_info > (3, 12), reason="segfault on 3.13")
99166
def test_wrap_generator():
100167
channel = []

0 commit comments

Comments
 (0)