|  | 
| 12 | 12 | # See the License for the specific language governing permissions and | 
| 13 | 13 | # limitations under the License. | 
| 14 | 14 | 
 | 
|  | 15 | +from __future__ import annotations | 
|  | 16 | + | 
| 15 | 17 | import logging | 
| 16 |  | -from typing import ContextManager, Optional, Tuple | 
|  | 18 | +from typing import TYPE_CHECKING, Optional, Tuple | 
| 17 | 19 | 
 | 
| 18 | 20 | from celery import registry  # pylint: disable=no-name-in-module | 
| 19 | 21 | from celery.app.task import Task | 
| 20 | 22 | 
 | 
| 21 | 23 | from opentelemetry.semconv.trace import SpanAttributes | 
| 22 | 24 | from opentelemetry.trace import Span | 
| 23 | 25 | 
 | 
|  | 26 | +if TYPE_CHECKING: | 
|  | 27 | +    from contextlib import AbstractContextManager | 
|  | 28 | + | 
| 24 | 29 | logger = logging.getLogger(__name__) | 
| 25 | 30 | 
 | 
| 26 | 31 | # Celery Context key | 
| @@ -123,7 +128,7 @@ def attach_context( | 
| 123 | 128 |     task: Optional[Task], | 
| 124 | 129 |     task_id: str, | 
| 125 | 130 |     span: Span, | 
| 126 |  | -    activation: ContextManager[Span], | 
|  | 131 | +    activation: AbstractContextManager[Span], | 
| 127 | 132 |     token: Optional[object], | 
| 128 | 133 |     is_publish: bool = False, | 
| 129 | 134 | ) -> None: | 
| @@ -171,7 +176,7 @@ def detach_context(task, task_id, is_publish=False) -> None: | 
| 171 | 176 | 
 | 
| 172 | 177 | def retrieve_context( | 
| 173 | 178 |     task, task_id, is_publish=False | 
| 174 |  | -) -> Optional[Tuple[Span, ContextManager[Span], Optional[object]]]: | 
|  | 179 | +) -> Optional[Tuple[Span, AbstractContextManager[Span], Optional[object]]]: | 
| 175 | 180 |     """Helper to retrieve an active `Span`, `ContextManager` and context token | 
| 176 | 181 |     stored in a `Task` instance | 
| 177 | 182 |     """ | 
|  | 
0 commit comments