1818import random
1919import threading
2020from contextlib import contextmanager
21- from typing import Iterator , Optional , Sequence , Tuple
21+ from types import TracebackType
22+ from typing import Iterator , Optional , Sequence , Tuple , Type
2223
2324from opentelemetry import trace as trace_api
2425from opentelemetry .context import Context
2526from opentelemetry .sdk import util
2627from opentelemetry .sdk .util import BoundedDict , BoundedList
2728from opentelemetry .trace import sampling
29+ from opentelemetry .trace .status import Status , StatusCanonicalCode , StatusError
2830from opentelemetry .util import time_ns , types
2931
3032logger = logging .getLogger (__name__ )
@@ -134,6 +136,7 @@ def __init__(
134136 links : Sequence [trace_api .Link ] = (),
135137 kind : trace_api .SpanKind = trace_api .SpanKind .INTERNAL ,
136138 span_processor : SpanProcessor = SpanProcessor (),
139+ auto_update_status : bool = True ,
137140 ) -> None :
138141
139142 self .name = name
@@ -143,9 +146,17 @@ def __init__(
143146 self .trace_config = trace_config
144147 self .resource = resource
145148 self .kind = kind
149+ self ._auto_update_status = auto_update_status
146150
147151 self .span_processor = span_processor
148- self .status = trace_api .Status ()
152+ # FIXME: check the statement below.
153+ # The status is a response code. A span that has been just initialized
154+ # should not have a response code yet. This also makes it possible to
155+ # detect if the span status has been manually set to some value, in
156+ # that case, exiting the span will not cause the span to have its
157+ # status set to some status if an exception was raised while the span
158+ # was active.
159+ self .status = None
149160 self ._lock = threading .Lock ()
150161
151162 if attributes is None :
@@ -174,7 +185,10 @@ def __repr__(self):
174185 )
175186
176187 def __str__ (self ):
177- return '{}(name="{}", context={}, kind={}, parent={}, start_time={}, end_time={})' .format (
188+ return (
189+ '{}(name="{}", context={}, kind={}, '
190+ "parent={}, start_time={}, end_time={})"
191+ ).format (
178192 type (self ).__name__ ,
179193 self .name ,
180194 self .context ,
@@ -275,6 +289,30 @@ def set_status(self, status: trace_api.Status) -> None:
275289 return
276290 self .status = status
277291
292+ def __exit__ (
293+ self ,
294+ exc_type : Optional [Type [BaseException ]],
295+ exc_val : Optional [BaseException ],
296+ exc_tb : Optional [TracebackType ],
297+ ) -> None :
298+ """Ends context manager and calls `end` on the `Span`."""
299+
300+ if self .status is None and self ._auto_update_status :
301+ if exc_val is not None :
302+
303+ if isinstance (exc_val , StatusError ):
304+ self .set_status (
305+ Status (exc_val .canonical_code , exc_val .description )
306+ )
307+
308+ else :
309+ self .set_status (Status (StatusCanonicalCode .UNKNOWN ))
310+
311+ else :
312+ self .set_status (Status (StatusCanonicalCode .OK ))
313+
314+ super (Span , self ).__exit__ (exc_type , exc_val , exc_tb )
315+
278316
279317def generate_span_id () -> int :
280318 """Get a new random span ID.
@@ -334,14 +372,15 @@ def start_as_current_span(
334372 span = self .start_span (name , parent , kind , attributes , links )
335373 return self .use_span (span , end_on_exit = True )
336374
337- def start_span (
375+ def start_span ( # pylint: disable=arguments-differ, too-many-locals
338376 self ,
339377 name : str ,
340378 parent : trace_api .ParentSpan = trace_api .Tracer .CURRENT_SPAN ,
341379 kind : trace_api .SpanKind = trace_api .SpanKind .INTERNAL ,
342380 attributes : Optional [types .Attributes ] = None ,
343381 links : Sequence [trace_api .Link ] = (),
344382 start_time : Optional [int ] = None ,
383+ auto_update_status : bool = True ,
345384 ) -> "Span" :
346385 """See `opentelemetry.trace.Tracer.start_span`."""
347386
@@ -401,6 +440,7 @@ def start_span(
401440 span_processor = self ._active_span_processor ,
402441 kind = kind ,
403442 links = links ,
443+ auto_update_status = auto_update_status ,
404444 )
405445 span .start (start_time = start_time )
406446 else :
0 commit comments