Skip to content

Commit df31a30

Browse files
sentrivanaclaude
andcommitted
ref: Use early returns instead of nullcontext where possible
Replace nullcontext + if-span-is-not-None patterns with early returns when the with block ends with a return. This removes unnecessary nesting and nullcontext imports from 23 files. The only remaining nullcontext usage is in pyreqwest.py where the function is a generator context manager that must yield. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 6272d36 commit df31a30

23 files changed

Lines changed: 679 additions & 736 deletions

File tree

sentry_sdk/integrations/aiomysql.py

Lines changed: 22 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
)
1515
from sentry_sdk.utils import (
1616
capture_internal_exceptions,
17-
nullcontext,
1817
parse_version,
1918
)
2019

@@ -180,36 +179,34 @@ async def _inner(self: "Connection") -> T:
180179
"sentry.origin": AioMySQLIntegration.origin,
181180
} | breadcrumb_data
182181

183-
span_ctx = nullcontext()
184-
if sentry_sdk.traces.get_current_span() is not None:
185-
span_ctx = sentry_sdk.traces.start_span(
186-
name="connect", attributes=span_attributes
187-
)
188-
with span_ctx as span:
182+
if sentry_sdk.traces.get_current_span() is None:
183+
return await f(self)
184+
185+
with sentry_sdk.traces.start_span(
186+
name="connect", attributes=span_attributes
187+
):
189188
with capture_internal_exceptions():
190189
sentry_sdk.add_breadcrumb(
191190
message="connect", category="query", data=breadcrumb_data
192191
)
193-
res = await f(self)
194-
else:
195-
connect_data = _get_connect_data(self)
196-
197-
with sentry_sdk.start_span(
198-
op=OP.DB,
199-
name="connect",
200-
origin=AioMySQLIntegration.origin,
201-
) as span:
202-
_set_db_data(span, self)
192+
return await f(self)
203193

204-
with capture_internal_exceptions():
205-
sentry_sdk.add_breadcrumb(
206-
message="connect",
207-
category="query",
208-
data=connect_data,
209-
)
210-
res = await f(self)
194+
connect_data = _get_connect_data(self)
211195

212-
return res
196+
with sentry_sdk.start_span(
197+
op=OP.DB,
198+
name="connect",
199+
origin=AioMySQLIntegration.origin,
200+
) as span:
201+
_set_db_data(span, self)
202+
203+
with capture_internal_exceptions():
204+
sentry_sdk.add_breadcrumb(
205+
message="connect",
206+
category="query",
207+
data=connect_data,
208+
)
209+
return await f(self)
213210

214211
return _inner
215212

sentry_sdk/integrations/arq.py

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
capture_internal_exceptions,
1515
ensure_integration_enabled,
1616
event_from_exception,
17-
nullcontext,
1817
parse_version,
1918
reraise,
2019
)
@@ -80,16 +79,16 @@ async def _sentry_enqueue_job(
8079
return await old_enqueue_job(self, function, *args, **kwargs)
8180

8281
if has_span_streaming_enabled(client.options):
83-
span_ctx = nullcontext()
84-
if sentry_sdk.traces.get_current_span() is not None:
85-
span_ctx = sentry_sdk.traces.start_span(
86-
name=function,
87-
attributes={
88-
"sentry.op": OP.QUEUE_SUBMIT_ARQ,
89-
"sentry.origin": ArqIntegration.origin,
90-
},
91-
)
92-
with span_ctx:
82+
if sentry_sdk.traces.get_current_span() is None:
83+
return await old_enqueue_job(self, function, *args, **kwargs)
84+
85+
with sentry_sdk.traces.start_span(
86+
name=function,
87+
attributes={
88+
"sentry.op": OP.QUEUE_SUBMIT_ARQ,
89+
"sentry.origin": ArqIntegration.origin,
90+
},
91+
):
9392
return await old_enqueue_job(self, function, *args, **kwargs)
9493

9594
with sentry_sdk.start_span(

sentry_sdk/integrations/asyncpg.py

Lines changed: 27 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
)
1717
from sentry_sdk.utils import (
1818
capture_internal_exceptions,
19-
nullcontext,
2019
parse_version,
2120
)
2221

@@ -236,42 +235,39 @@ async def _inner(*args: "Any", **kwargs: "Any") -> "T":
236235
except IndexError:
237236
pass
238237

239-
span_ctx = nullcontext()
240-
if sentry_sdk.traces.get_current_span() is not None:
241-
span_ctx = sentry_sdk.traces.start_span(
242-
name="connect", attributes=span_attributes
243-
)
244-
with span_ctx as span:
238+
if sentry_sdk.traces.get_current_span() is None:
239+
return await f(*args, **kwargs)
240+
241+
with sentry_sdk.traces.start_span(
242+
name="connect", attributes=span_attributes
243+
):
245244
with capture_internal_exceptions():
246245
sentry_sdk.add_breadcrumb(
247246
message="connect", category="query", data=span_attributes
248247
)
249-
res = await f(*args, **kwargs)
248+
return await f(*args, **kwargs)
250249

251-
else:
252-
with sentry_sdk.start_span(
253-
op=OP.DB,
254-
name="connect",
255-
origin=AsyncPGIntegration.origin,
256-
) as span:
257-
span.set_data(SPANDATA.DB_SYSTEM, "postgresql")
258-
if addr:
259-
try:
260-
span.set_data(SPANDATA.SERVER_ADDRESS, addr[0])
261-
span.set_data(SPANDATA.SERVER_PORT, addr[1])
262-
except IndexError:
263-
pass
264-
span.set_data(SPANDATA.DB_NAME, database)
265-
span.set_data(SPANDATA.DB_USER, user)
266-
span.set_data(SPANDATA.DB_DRIVER_NAME, "asyncpg")
267-
268-
with capture_internal_exceptions():
269-
sentry_sdk.add_breadcrumb(
270-
message="connect", category="query", data=span._data
271-
)
272-
res = await f(*args, **kwargs)
250+
with sentry_sdk.start_span(
251+
op=OP.DB,
252+
name="connect",
253+
origin=AsyncPGIntegration.origin,
254+
) as span:
255+
span.set_data(SPANDATA.DB_SYSTEM, "postgresql")
256+
if addr:
257+
try:
258+
span.set_data(SPANDATA.SERVER_ADDRESS, addr[0])
259+
span.set_data(SPANDATA.SERVER_PORT, addr[1])
260+
except IndexError:
261+
pass
262+
span.set_data(SPANDATA.DB_NAME, database)
263+
span.set_data(SPANDATA.DB_USER, user)
264+
span.set_data(SPANDATA.DB_DRIVER_NAME, "asyncpg")
273265

274-
return res
266+
with capture_internal_exceptions():
267+
sentry_sdk.add_breadcrumb(
268+
message="connect", category="query", data=span._data
269+
)
270+
return await f(*args, **kwargs)
275271

276272
return _inner
277273

sentry_sdk/integrations/celery/__init__.py

Lines changed: 44 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
SENSITIVE_DATA_SUBSTITUTE,
2424
capture_internal_exceptions,
2525
event_from_exception,
26-
nullcontext,
2726
reraise,
2827
)
2928

@@ -420,17 +419,17 @@ def _inner(*args: "Any", **kwargs: "Any") -> "Any":
420419
span_streaming = has_span_streaming_enabled(client.options)
421420

422421
try:
423-
span_ctx: "Any"
422+
if span_streaming and get_current_span() is None:
423+
return f(*args, **kwargs)
424+
424425
if span_streaming:
425-
span_ctx = nullcontext()
426-
if get_current_span() is not None:
427-
span_ctx = sentry_sdk.traces.start_span(
428-
name=task.name,
429-
attributes={
430-
"sentry.op": OP.QUEUE_PROCESS,
431-
"sentry.origin": CeleryIntegration.origin,
432-
},
433-
)
426+
span_ctx = sentry_sdk.traces.start_span(
427+
name=task.name,
428+
attributes={
429+
"sentry.op": OP.QUEUE_PROCESS,
430+
"sentry.origin": CeleryIntegration.origin,
431+
},
432+
)
434433
else:
435434
span_ctx = sentry_sdk.start_span(
436435
op=OP.QUEUE_PROCESS,
@@ -439,44 +438,42 @@ def _inner(*args: "Any", **kwargs: "Any") -> "Any":
439438
)
440439

441440
with span_ctx as span:
442-
if span is not None:
443-
if isinstance(span, StreamedSpan):
444-
set_on_span = span.set_attribute
445-
else:
446-
set_on_span = span.set_data
447-
448-
_set_messaging_destination_name(task, span)
449-
450-
latency = None
451-
with capture_internal_exceptions():
452-
if (
453-
task.request.headers is not None
454-
and "sentry-task-enqueued-time" in task.request.headers
455-
):
456-
latency = (
457-
_now_seconds_since_epoch()
458-
- task.request.headers.pop("sentry-task-enqueued-time")
459-
)
460-
461-
if latency is not None:
462-
latency *= 1000 # milliseconds
463-
set_on_span(SPANDATA.MESSAGING_MESSAGE_RECEIVE_LATENCY, latency)
464-
465-
with capture_internal_exceptions():
466-
set_on_span(SPANDATA.MESSAGING_MESSAGE_ID, task.request.id)
467-
468-
with capture_internal_exceptions():
469-
set_on_span(
470-
SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT,
471-
task.request.retries,
441+
if isinstance(span, StreamedSpan):
442+
set_on_span = span.set_attribute
443+
else:
444+
set_on_span = span.set_data
445+
446+
_set_messaging_destination_name(task, span)
447+
448+
latency = None
449+
with capture_internal_exceptions():
450+
if (
451+
task.request.headers is not None
452+
and "sentry-task-enqueued-time" in task.request.headers
453+
):
454+
latency = _now_seconds_since_epoch() - task.request.headers.pop(
455+
"sentry-task-enqueued-time"
472456
)
473457

474-
with capture_internal_exceptions():
475-
with task.app.connection() as conn:
476-
set_on_span(
477-
SPANDATA.MESSAGING_SYSTEM,
478-
conn.transport.driver_type,
479-
)
458+
if latency is not None:
459+
latency *= 1000 # milliseconds
460+
set_on_span(SPANDATA.MESSAGING_MESSAGE_RECEIVE_LATENCY, latency)
461+
462+
with capture_internal_exceptions():
463+
set_on_span(SPANDATA.MESSAGING_MESSAGE_ID, task.request.id)
464+
465+
with capture_internal_exceptions():
466+
set_on_span(
467+
SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT,
468+
task.request.retries,
469+
)
470+
471+
with capture_internal_exceptions():
472+
with task.app.connection() as conn:
473+
set_on_span(
474+
SPANDATA.MESSAGING_SYSTEM,
475+
conn.transport.driver_type,
476+
)
480477

481478
return f(*args, **kwargs)
482479
except Exception:

sentry_sdk/integrations/cohere.py

Lines changed: 24 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
from sentry_sdk.utils import (
2121
capture_internal_exceptions,
2222
event_from_exception,
23-
nullcontext,
2423
reraise,
2524
)
2625

@@ -257,16 +256,15 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any":
257256
)
258257

259258
if is_span_streaming_enabled:
260-
if sentry_sdk.traces.get_current_span() is not None:
261-
span_ctx = sentry_sdk.traces.start_span(
262-
name="Cohere Embedding Creation",
263-
attributes={
264-
"sentry.op": consts.OP.COHERE_EMBEDDINGS_CREATE,
265-
"sentry.origin": CohereIntegration.origin,
266-
},
267-
)
268-
else:
269-
span_ctx = nullcontext()
259+
if sentry_sdk.traces.get_current_span() is None:
260+
return f(*args, **kwargs)
261+
span_ctx = sentry_sdk.traces.start_span(
262+
name="Cohere Embedding Creation",
263+
attributes={
264+
"sentry.op": consts.OP.COHERE_EMBEDDINGS_CREATE,
265+
"sentry.origin": CohereIntegration.origin,
266+
},
267+
)
270268
else:
271269
span_ctx = get_start_span_function()(
272270
op=consts.OP.COHERE_EMBEDDINGS_CREATE,
@@ -275,31 +273,30 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any":
275273
)
276274

277275
with span_ctx as span:
278-
if span is not None:
279-
if "texts" in kwargs and (
280-
should_send_default_pii() and integration.include_prompts
276+
if "texts" in kwargs and (
277+
should_send_default_pii() and integration.include_prompts
278+
):
279+
if isinstance(kwargs["texts"], str):
280+
set_data_normalized(span, SPANDATA.AI_TEXTS, [kwargs["texts"]])
281+
elif (
282+
isinstance(kwargs["texts"], list)
283+
and len(kwargs["texts"]) > 0
284+
and isinstance(kwargs["texts"][0], str)
281285
):
282-
if isinstance(kwargs["texts"], str):
283-
set_data_normalized(span, SPANDATA.AI_TEXTS, [kwargs["texts"]])
284-
elif (
285-
isinstance(kwargs["texts"], list)
286-
and len(kwargs["texts"]) > 0
287-
and isinstance(kwargs["texts"][0], str)
288-
):
289-
set_data_normalized(
290-
span, SPANDATA.AI_INPUT_MESSAGES, kwargs["texts"]
291-
)
286+
set_data_normalized(
287+
span, SPANDATA.AI_INPUT_MESSAGES, kwargs["texts"]
288+
)
292289

293-
if "model" in kwargs:
294-
set_data_normalized(span, SPANDATA.AI_MODEL_ID, kwargs["model"])
290+
if "model" in kwargs:
291+
set_data_normalized(span, SPANDATA.AI_MODEL_ID, kwargs["model"])
295292
try:
296293
res = f(*args, **kwargs)
297294
except Exception as e:
298295
exc_info = sys.exc_info()
299296
with capture_internal_exceptions():
300297
_capture_exception(e)
301298
reraise(*exc_info)
302-
if span is not None and (
299+
if (
303300
hasattr(res, "meta")
304301
and hasattr(res.meta, "billed_units")
305302
and hasattr(res.meta.billed_units, "input_tokens")

0 commit comments

Comments
 (0)