From a0b5f231b591a983838a79fe43b191dd266e96f7 Mon Sep 17 00:00:00 2001
From: hyoinandout <hobbyhyoin@naver.com>
Date: Wed, 4 Sep 2024 15:07:52 +0900
Subject: [PATCH] Add type annotations

---
 .../src/opentelemetry/context/__init__.py            |  7 +++----
 .../src/opentelemetry/context/context.py             |  5 +++--
 .../src/opentelemetry/context/contextvars_context.py | 12 +++++-------
 opentelemetry-api/tests/context/base_context.py      |  2 +-
 4 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/opentelemetry-api/src/opentelemetry/context/__init__.py b/opentelemetry-api/src/opentelemetry/context/__init__.py
index cc07155bd35..88a9b9f4ac0 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
 
@@ -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")
 
diff --git a/opentelemetry-api/src/opentelemetry/context/context.py b/opentelemetry-api/src/opentelemetry/context/context.py
index 518f09f2b8f..1214b62e166 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 9136f09df0c..8442ca733bf 100644
--- a/opentelemetry-api/src/opentelemetry/context/contextvars_context.py
+++ b/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"]
diff --git a/opentelemetry-api/tests/context/base_context.py b/opentelemetry-api/tests/context/base_context.py
index 05acc95d89a..5c628225fce 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("some garbage")  # type:ignore
 
         def test_detach_out_of_order(self):
             t1 = context.attach(context.set_value("c", 1))