Add type annotations

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
 
@@ -130,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.
 
@@ -143,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.
 
@@ -152,8 +153,6 @@ def detach(token: object) -> None:
     """
     try:
         _RUNTIME_CONTEXT.detach(token)
-    except TypeError:
-        logger.exception("Expected an instance of Token, got None")
     except Exception:  # pylint: disable=broad-exception-caught
         logger.exception("Failed to detach context")
 

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:

opentelemetry-api/src/opentelemetry/context/contextvars_context.py

@@ -11,7 +11,8 @@
 # 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
+import typing
+from contextvars import ContextVar, Token
 
 from opentelemetry.context.context import Context, _RuntimeContext
 
@@ -28,7 +29,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,16 +42,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.
         """
-        try:
-            self._current_context.reset(token)
-        except TypeError:
-            raise TypeError
+        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("some garbage")  # type:ignore
 
         def test_detach_out_of_order(self):
             t1 = context.attach(context.set_value("c", 1))