diff --git a/CHANGELOG.md b/CHANGELOG.md index e71d5a7b6a..2bbb075fbb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased +- Add type annotations to context's attach & detach + ([#4164](https://github.com/open-telemetry/opentelemetry-python/pull/4346)) - Add `attributes` field in `metrics.get_meter` wrapper function ([#4364](https://github.com/open-telemetry/opentelemetry-python/pull/4364)) - Add Python 3.13 support diff --git a/opentelemetry-api/src/opentelemetry/context/__init__.py b/opentelemetry-api/src/opentelemetry/context/__init__.py index b9a5a17b2f..6e3ac9531b 100644 --- a/opentelemetry-api/src/opentelemetry/context/__init__.py +++ b/opentelemetry-api/src/opentelemetry/context/__init__.py @@ -14,6 +14,7 @@ import logging import typing +from contextvars import Token from os import environ from uuid import uuid4 @@ -128,7 +129,7 @@ def get_current() -> Context: return _RUNTIME_CONTEXT.get_current() -def attach(context: Context) -> object: +def attach(context: Context) -> Token[Context]: """Associates a Context with the caller's current execution unit. Returns a token that can be used to restore the previous Context. @@ -141,7 +142,7 @@ def attach(context: Context) -> object: return _RUNTIME_CONTEXT.attach(context) -def detach(token: object) -> None: +def detach(token: Token[Context]) -> None: """Resets the Context associated with the caller's current execution unit to the value it had before attaching a specified Context. diff --git a/opentelemetry-api/src/opentelemetry/context/context.py b/opentelemetry-api/src/opentelemetry/context/context.py index 518f09f2b8..1214b62e16 100644 --- a/opentelemetry-api/src/opentelemetry/context/context.py +++ b/opentelemetry-api/src/opentelemetry/context/context.py @@ -14,6 +14,7 @@ import typing from abc import ABC, abstractmethod +from contextvars import Token class Context(typing.Dict[str, object]): @@ -29,7 +30,7 @@ class _RuntimeContext(ABC): """ @abstractmethod - def attach(self, context: Context) -> object: + def attach(self, context: Context) -> Token[Context]: """Sets the current `Context` object. Returns a token that can be used to reset to the previous `Context`. @@ -42,7 +43,7 @@ def get_current(self) -> Context: """Returns the current `Context` object.""" @abstractmethod - def detach(self, token: object) -> None: + def detach(self, token: Token[Context]) -> None: """Resets Context to a previous value Args: diff --git a/opentelemetry-api/src/opentelemetry/context/contextvars_context.py b/opentelemetry-api/src/opentelemetry/context/contextvars_context.py index 5f606764fc..9479839fe5 100644 --- a/opentelemetry-api/src/opentelemetry/context/contextvars_context.py +++ b/opentelemetry-api/src/opentelemetry/context/contextvars_context.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from contextvars import ContextVar +from contextvars import ContextVar, Token from opentelemetry.context.context import Context, _RuntimeContext @@ -28,7 +28,7 @@ def __init__(self) -> None: self._CONTEXT_KEY, default=Context() ) - def attach(self, context: Context) -> object: + def attach(self, context: Context) -> Token[Context]: """Sets the current `Context` object. Returns a token that can be used to reset to the previous `Context`. @@ -41,13 +41,13 @@ def get_current(self) -> Context: """Returns the current `Context` object.""" return self._current_context.get() - def detach(self, token: object) -> None: + def detach(self, token: Token[Context]) -> None: """Resets Context to a previous value Args: token: A reference to a previous Context. """ - self._current_context.reset(token) # type: ignore + self._current_context.reset(token) __all__ = ["ContextVarsRuntimeContext"]