Coverage for lasso/utils/rich_progress_bars.py: 42%

1import math 

2import time 

3from typing import Any 

4 

5from rich.progress import ProgressColumn 

6 

7 

8class PlaceHolderBar: 

9 """This bar is simply a placeholder bar""" 

10 

11 finished: bool 

12 tasks: list = [] 

13 

14 # noinspection PyUnusedLocal 

15 # pylint: disable = unused-argument 

16 def __init__(self, **kwargs): 

17 """This is a placeholder to not clutter console during testing""" 

18 self.finished: False 

19 

20 # noinspection PyUnusedLocal 

21 # pylint: disable = unused-argument 

22 def render(self, task: Any) -> str: 

23 """returns the planned output: empty string""" 

24 return "" 

25 

26 def add_task(self, description: str, total: int) -> int: 

27 """Adds a new task""" 

28 self.tasks.append([description, total, 0]) 

29 # entry in list is tuple of description, total tasks, remaining tasks 

30 return len(self.tasks) - 1 

31 

32 def advance(self, task_id): 

33 """advances the given task""" 

34 prog = self.tasks[task_id][2] 

35 prog += 1 

36 self.tasks[task_id][2] = prog 

37 if prog == self.tasks[task_id][1]: 

38 self.finished = True 

39 

40 def __enter__(self): 

41 self.finished = False 

42 

43 def __exit__(self, exception_type, exception_value, traceback): 

44 self.finished = True 

45 

46 

47class WorkingDots(ProgressColumn): 

48 """TODO: add description""" 

49 

50 max_refresh = 0.5 

51 is_silenced: bool = False 

52 

53 def __init__(self, output=True): 

54 self.counter = 0 

55 if not output: 

56 self.is_silenced = True 

57 super().__init__() 

58 

59 def render(self, task: Any) -> str: 

60 self.counter += 1 

61 if self.is_silenced: 

62 return "" 

63 if task.completed == task.total: 

64 msg = "..." 

65 elif self.counter % 3 == 0: 

66 msg = ". " 

67 elif self.counter % 3 == 1: 

68 msg = ".. " 

69 else: 

70 msg = "..." 

71 self.counter = 2 

72 return msg 

73 

74 

75class SubsamplingWaitTime(ProgressColumn): 

76 """TODO: add description""" 

77 

78 max_refresh = 0.5 

79 

80 def __init__(self, n_proc: int): 

81 super().__init__() 

82 

83 # Last time we updated 

84 self.last_time = time.time() 

85 # Cumulative time of all completed sub-sampling processes 

86 self.cum_time = 0 

87 # Number of parallel running processes 

88 self.n_proc = n_proc 

89 self.t_rem = -1 

90 

91 def render(self, task: Any) -> str: 

92 """TODO: add description?""" 

93 

94 if task.completed == task.total: 

95 return "Time remaining: 00:00" 

96 

97 if self.cum_time > 0: 

98 avrg_time = self.cum_time / max(1, task.completed) 

99 rem_tasks = task.total - task.completed 

100 gr_tasks = math.floor(rem_tasks / self.n_proc) 

101 if (rem_tasks % self.n_proc) != 0: 

102 gr_tasks += 1 

103 

104 total_time = gr_tasks * avrg_time 

105 if self.t_rem < 0 or self.t_rem > total_time: 

106 self.t_rem = total_time 

107 

108 t_out = self.t_rem - (time.time() - self.last_time) 

109 mins = str(math.floor(t_out / 60)) 

110 secs = str(math.trunc(t_out % 60)) 

111 

112 if len(mins) == 1: 

113 mins = "0" + mins 

114 if len(secs) == 1: 

115 secs = "0" + secs 

116 out_str = "Time remaining: " + mins + ":" + secs 

117 return out_str 

118 

119 return "Time remaining: --:--" 

120 

121 def update_avrg(self, new_time: float): 

122 """TODO: add description""" 

123 

124 self.cum_time += new_time 

125 self.last_time = time.time()