|
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