Coverage for bioimageio/spec/_internal/progress.py: 100%

1from typing import Any, Optional, Protocol 

2 

3from rich.progress import Progress 

4 

5 

6class Progressbar(Protocol): 

7 """Progressbar protocol modeled after tqdm""" 

8 

9 total: Optional[int] 

10 

11 def update(self, increment: int, /) -> Any: ... 

12 

13 def reset(self): ... 

14 

15 def close(self): ... 

16 

17 def set_description(self, description: str, /, refresh: bool = True): ... 

18 

19 

20class RichTaskBar: 

21 def __init__(self, description: str, *, parent: Progress, total: Optional[int]): 

22 super().__init__() 

23 self.task_id = parent.add_task(description, total=total) 

24 self.parent = parent 

25 self.total = total 

26 

27 def update(self, increment: int, /): 

28 self.parent.advance(self.task_id, increment) 

29 

30 def reset(self): 

31 self.parent.reset(self.task_id) 

32 

33 def close(self): 

34 self.parent.remove_task(self.task_id) 

35 

36 def set_description(self, description: str, /, refresh: bool = True): 

37 self.parent.update(self.task_id, description=description, refresh=refresh) 

38 

39 

40class RichOverallProgress: 

41 def __init__(self): 

42 super().__init__() 

43 self.progress = Progress() 

44 

45 def __call__(self, description: str = "", *, total: Optional[int] = None): 

46 return RichTaskBar(description, parent=self.progress, total=total)