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

1from abc import abstractmethod 

2from typing import Any, Optional, Protocol 

3 

4from rich.progress import Progress 

5 

6 

7class ProgressbarLike(Protocol): 

8 """Progressbar protocol modeled after tqdm""" 

9 

10 total: Optional[int] 

11 

12 @abstractmethod 

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

14 

15 @abstractmethod 

16 def reset(self): ... 

17 

18 @abstractmethod 

19 def close(self): ... 

20 

21 @abstractmethod 

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

23 

24 

25class RichTaskBar(ProgressbarLike): 

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

27 super().__init__() 

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

29 self.parent = parent 

30 self.total = total 

31 

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

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

34 

35 def reset(self): 

36 self.parent.reset(self.task_id) 

37 

38 def close(self): 

39 self.parent.remove_task(self.task_id) 

40 

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

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

43 

44 

45class RichOverallProgress: 

46 def __init__(self): 

47 super().__init__() 

48 self.progress = Progress() 

49 

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

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