diff --git a/CHANGELOG.md b/CHANGELOG.md index e194be58ee6..7e0fea3d903 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 + ([#4346](https://github.com/open-telemetry/opentelemetry-python/pull/4346)) - Fix OTLP encoders missing instrumentation scope schema url and attributes ([#4359](https://github.com/open-telemetry/opentelemetry-python/pull/4359)) - prometheus-exporter: fix labels out of place for data points with different diff --git a/docs/conf.py b/docs/conf.py index 965a806d06a..5e8037488bf 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -150,6 +150,10 @@ "py:class", "opentelemetry.sdk.metrics._internal.aggregation._Aggregation", ), + ( + "py:class", + "_contextvars.Token", + ), ] # Add any paths that contain templates here, relative to this directory. diff --git a/opentelemetry-api/src/opentelemetry/context/__init__.py b/opentelemetry-api/src/opentelemetry/context/__init__.py index b9a5a17b2f1..cad7f951428 100644 --- a/opentelemetry-api/src/opentelemetry/context/__init__.py +++ b/opentelemetry-api/src/opentelemetry/context/__init__.py @@ -12,8 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import logging import typing +from contextvars import Token from os import environ from uuid import uuid4 @@ -128,7 +131,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 +144,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 518f09f2b8f..c1ef9cfbb6b 100644 --- a/opentelemetry-api/src/opentelemetry/context/context.py +++ b/opentelemetry-api/src/opentelemetry/context/context.py @@ -12,8 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import typing from abc import ABC, abstractmethod +from contextvars import Token class Context(typing.Dict[str, object]): @@ -29,7 +32,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 +45,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 5f606764fcc..dceee263482 100644 --- a/opentelemetry-api/src/opentelemetry/context/contextvars_context.py +++ b/opentelemetry-api/src/opentelemetry/context/contextvars_context.py @@ -11,7 +11,10 @@ # 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 __future__ import annotations + +from contextvars import ContextVar, Token from opentelemetry.context.context import Context, _RuntimeContext @@ -28,7 +31,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 +44,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"] diff --git a/opentelemetry-api/tests/context/base_context.py b/opentelemetry-api/tests/context/base_context.py index 05acc95d89a..395229b5208 100644 --- a/opentelemetry-api/tests/context/base_context.py +++ b/opentelemetry-api/tests/context/base_context.py @@ -64,7 +64,7 @@ def test_attach(self): self.assertEqual("yyy", context.get_value("a")) with self.assertLogs(level=ERROR): - context.detach("some garbage") + context.detach(token) def test_detach_out_of_order(self): t1 = context.attach(context.set_value("c", 1))