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

1from __future__ import annotations 

2 

3from abc import abstractmethod 

4from typing import Any, Protocol 

5 

6from rich.progress import Progress 

7 

8 

9class ProgressbarLike(Protocol): 

10 """Progressbar protocol modeled after tqdm""" 

11 

12 total: int | None 

13 

14 @abstractmethod 

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

16 

17 @abstractmethod 

18 def reset(self): ... 

19 

20 @abstractmethod 

21 def close(self): ... 

22 

23 @abstractmethod 

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

25 

26 

27class RichTaskBar(ProgressbarLike): 

28 def __init__(self, description: str, *, parent: Progress, total: int | None): 

29 super().__init__() 

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

31 self.parent = parent 

32 self.total = total 

33 

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

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

36 

37 def reset(self): 

38 self.parent.reset(self.task_id) 

39 

40 def close(self): 

41 self.parent.remove_task(self.task_id) 

42 

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

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

45 

46 

47class RichOverallProgress: 

48 def __init__(self): 

49 super().__init__() 

50 self.progress = Progress() 

51 

52 def __call__(self, description: str = "", *, total: int | None = None): 

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