codekingpro/portable-devtools

1# Copyright The OpenTelemetry Authors2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from typing import Optional15 16from opentelemetry.context import create_key, get_value, set_value17from opentelemetry.context.context import Context18from opentelemetry.trace.span import INVALID_SPAN, Span19 20SPAN_KEY = "current-span"21_SPAN_KEY = create_key("current-span")22 23 24def set_span_in_context(25    span: Span, context: Optional[Context] = None26) -> Context:27    """Set the span in the given context.28 29    Args:30        span: The Span to set.31        context: a Context object. if one is not passed, the32            default current context is used instead.33    """34    ctx = set_value(_SPAN_KEY, span, context=context)35    return ctx36 37 38def get_current_span(context: Optional[Context] = None) -> Span:39    """Retrieve the current span.40 41    Args:42        context: A Context object. If one is not passed, the43            default current context is used instead.44 45    Returns:46        The Span set in the context if it exists. INVALID_SPAN otherwise.47    """48    span = get_value(_SPAN_KEY, context=context)49    if span is None or not isinstance(span, Span):50        return INVALID_SPAN51    return span52