Add typing information for contexts, token (#4346)

* Handle TypeError from opentelemetry.context.contextvars_context in detach

* Add type annotations

* In changelog update PR number

* remove unnecessary type:ignore

* fix tox -e mypy tests

* Apply suggestions from code review

Co-authored-by: Emídio Neto <[email protected]>

* Apply suggestions from code review

Co-authored-by: Emídio Neto <[email protected]>

* Update conf.py

* Update docs/conf.py

* Update docs/conf.py

---------

Co-authored-by: hyoinandout <[email protected]>
Co-authored-by: Leighton Chen <[email protected]>
Co-authored-by: Riccardo Magliocchetti <[email protected]>
Co-authored-by: Emídio Neto <[email protected]>

Author: 5 people

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

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.

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.
 

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:

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"]

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))