Coverage for /usr/local/lib/python3.12/site-packages/prefect/utilities/context.py: 29%

1from collections.abc import Generator 1 ctx1a

2from contextlib import contextmanager 1 ctx1a

3from contextvars import Context, ContextVar, Token 1 ctx1a

4from typing import TYPE_CHECKING, Any, Optional, cast 1 ctx1a

5from uuid import UUID 1 ctx1a

6 

7if TYPE_CHECKING: 7 ↛ 8line 7 didn't jump to line 8 because the condition on line 7 was never true1 ctx1a

8 from prefect.client.schemas.objects import FlowRun, TaskRun 

9 

10 

11@contextmanager 1 ctx1a

12def temporary_context(context: Context) -> Generator[None, Any, None]: 1 ctx1a

13 tokens: dict[ContextVar[Any], Token[Any]] = {} 

14 for key, value in context.items(): 

15 token = key.set(value) 

16 tokens[key] = token 

17 try: 

18 yield 

19 finally: 

20 for key, token in tokens.items(): 

21 key.reset(token) 

22 

23 

24def get_task_run_id() -> Optional[UUID]: 1 ctx1a

25 from prefect.context import TaskRunContext 

26 

27 context = TaskRunContext.get() 

28 if task_run := cast(Optional["TaskRun"], getattr(context, "task_run", None)): 

29 return task_run.id 

30 return None 

31 

32 

33def get_flow_run_id() -> Optional[UUID]: 1 ctx1a

34 from prefect.context import FlowRunContext 

35 

36 context = FlowRunContext.get() 

37 if flow_run := cast(Optional["FlowRun"], getattr(context, "flow_run", None)): 

38 return flow_run.id 

39 return None 

40 

41 

42def get_task_and_flow_run_ids() -> tuple[Optional[UUID], Optional[UUID]]: 1 ctx1a

43 """ 

44 Get the task run and flow run ids from the context, if available. 

45 

46 Returns: 

47 tuple[Optional[UUID], Optional[UUID]]: a tuple of the task run id and flow run id 

48 """ 

49 return get_task_run_id(), get_flow_run_id()