Coverage for adrs/data/progress.py: 67%

1from __future__ import annotations 

2 

3import sys 

4from contextlib import asynccontextmanager 

5from contextvars import ContextVar 

6from dataclasses import dataclass 

7from typing import AsyncIterator 

8 

9from rich.progress import ( 

10 BarColumn, 

11 Progress, 

12 SpinnerColumn, 

13 TaskID, 

14 TaskProgressColumn, 

15 TextColumn, 

16 TimeRemainingColumn, 

17) 

18 

19from adrs.console import console 

20 

21 

22_progress_var: ContextVar["Progress | None"] = ContextVar("adrs_progress", default=None) 

23 

24 

25@dataclass 

26class _NoOpBar: 

27 completed: int = 0 

28 total: int | None = None 

29 

30 def advance(self, n: int) -> None: 

31 return None 

32 

33 

34@dataclass 

35class _InnerBar: 

36 progress: Progress 

37 task_id: TaskID 

38 completed: int = 0 

39 total: int | None = None 

40 

41 def advance(self, n: int) -> None: 

42 self.progress.advance(self.task_id, n) 

43 self.completed += n 

44 

45 

46@dataclass 

47class _OuterBar: 

48 progress: Progress 

49 task_id: TaskID 

50 completed: int = 0 

51 total: int = 0 

52 

53 def advance(self, n: int = 1) -> None: 

54 self.progress.advance(self.task_id, n) 

55 self.completed += n 

56 

57 

58@asynccontextmanager 

59async def inner_task( 

60 description: str, total: int | None 

61) -> AsyncIterator[_NoOpBar | _InnerBar]: 

62 progress = _progress_var.get() 

63 if progress is None: 

64 yield _NoOpBar(completed=0, total=None) 

65 return 

66 

67 task_id = progress.add_task(description, total=total) 

68 bar = _InnerBar(progress=progress, task_id=task_id, completed=0, total=total) 

69 try: 

70 yield bar 

71 finally: 

72 progress.remove_task(task_id) 

73 

74 

75def _make_progress() -> Progress: 

76 return Progress( 

77 SpinnerColumn(), 

78 TextColumn("[progress.description]{task.description}"), 

79 BarColumn(), 

80 TaskProgressColumn(), 

81 TimeRemainingColumn(), 

82 console=console, 

83 disable=not sys.stdout.isatty(), 

84 ) 

85 

86 

87@asynccontextmanager 

88async def progress_context(total_topics: int) -> AsyncIterator[_OuterBar]: 

89 progress = _make_progress() 

90 progress.start() 

91 token = _progress_var.set(progress) 

92 outer_id = progress.add_task("topics", total=total_topics) 

93 outer = _OuterBar( 

94 progress=progress, task_id=outer_id, completed=0, total=total_topics 

95 ) 

96 try: 

97 yield outer 

98 finally: 

99 _progress_var.reset(token) 

100 progress.stop()