Coverage for src/threadful/bonus.py: 100%

49 statements  

« prev     ^ index     » next       coverage.py v7.4.3, created at 2024-11-19 20:26 +0100

1""" 

2Fun little util features. 

3 

4This module provides utility functions for cursor management and thread animation. 

5 

6Attributes: 

7 _print_kwargs (dict): A dictionary of keyword arguments for the print function. 

8""" 

9 

10import atexit 

11import sys 

12import time 

13import typing 

14from contextlib import contextmanager 

15 

16from .core import ThreadWithReturn 

17from .core import thread as threadify 

18 

19T = typing.TypeVar("T") 

20_print_kwargs: dict[str, typing.Any] = dict(file=sys.stderr, flush=True, end="\r", sep="") 

21 

22 

23# https://stackoverflow.com/questions/5174810/how-to-turn-off-blinking-cursor-in-command-window 

24 

25 

26def hide_cursor() -> None: 

27 """ 

28 Hides the cursor in the terminal. 

29 """ 

30 print("\033[?25l", end="", flush=True) 

31 atexit.register(show_cursor) # clean up when the script ends 

32 

33 

34def show_cursor() -> None: 

35 """ 

36 Shows the cursor in the terminal. 

37 """ 

38 print("\033[?25h", end="", flush=True) 

39 atexit.unregister(show_cursor) # clean up no longer required 

40 

41 

42@contextmanager 

43def toggle_cursor(enabled: bool = True) -> typing.Generator[None, None, None]: 

44 """ 

45 Toggles the visibility of the cursor in the terminal. 

46 

47 Args: 

48 enabled (bool): If True, the cursor is shown, otherwise it is hidden. 

49 """ 

50 if not enabled: 

51 yield 

52 return 

53 

54 hide_cursor() 

55 yield 

56 show_cursor() 

57 

58 

59T_Text: typing.TypeAlias = str | typing.Callable[[], str] 

60 

61 

62@threadify 

63def _animate_threaded( 

64 thread: ThreadWithReturn[T], 

65 text: T_Text = "", 

66 speed: float = 0.05, 

67 animation: tuple[str, ...] = ("⣷", "⣯", "⣟", "⡿", "⢿", "⣻", "⣽", "⣾"), 

68 clear_with: typing.Optional[str] = None, 

69) -> T: 

70 return _animate(thread, text=text, speed=speed, animation=animation, clear_with=clear_with) 

71 

72 

73def _animate( 

74 thread: ThreadWithReturn[T], 

75 text: T_Text = "", 

76 speed: float = 0.05, 

77 animation: tuple[str, ...] = ("⣷", "⣯", "⣟", "⡿", "⢿", "⣻", "⣽", "⣾"), 

78 clear_with: typing.Optional[str] = None, 

79) -> T: 

80 """ 

81 Private function to animate a loading spinner while a thread is running. 

82 

83 Args: 

84 thread (ThreadWithReturn): The thread to animate. 

85 speed (float): The speed of the animation. 

86 animation (tuple): The frames of the animation. 

87 clear_with (str): replace the animation with a specific character instead of clearing the text line 

88 

89 Returns: 

90 T: The result of the thread. 

91 """ 

92 idx = 0 

93 while not thread.is_done(): 

94 idx += 1 

95 _text = text() if callable(text) else text 

96 print(animation[idx % len(animation)], " ", _text, **_print_kwargs) 

97 time.sleep(speed) 

98 

99 # print enough spaces to clear text: 

100 _text = text() if callable(text) else text 

101 

102 if clear_with: 

103 print(clear_with, " ", _text, "\n", **_print_kwargs) 

104 else: 

105 buffer_spaces = len(_text) + 1 

106 print("\r ", " " * buffer_spaces, **_print_kwargs) 

107 

108 return thread.join() 

109 

110 

111@typing.overload 

112def animate( 

113 thread: ThreadWithReturn[T], 

114 threaded: typing.Literal[True], 

115 text: T_Text = "", 

116 speed: float = 0.05, 

117 animation: tuple[str, ...] = (), 

118 clear_with: typing.Optional[str] = None, 

119 _hide_cursor: bool = True, 

120) -> ThreadWithReturn[T]: 

121 """ 

122 Pass threaded=True to also thread the loading animation, clearing up the thread. 

123 """ 

124 

125 

126@typing.overload 

127def animate( 

128 thread: ThreadWithReturn[T], 

129 threaded: typing.Literal[False] = False, 

130 text: T_Text = "", 

131 speed: float = 0.05, 

132 animation: tuple[str, ...] = (), 

133 clear_with: typing.Optional[str] = None, 

134 _hide_cursor: bool = True, 

135) -> T: 

136 """ 

137 Default behavior: run the animation sync. 

138 """ 

139 

140 

141def animate( 

142 thread: ThreadWithReturn[T], 

143 threaded: bool = False, 

144 text: T_Text = "", 

145 speed: float = 0.05, 

146 animation: tuple[str, ...] = ("⣷", "⣯", "⣟", "⡿", "⢿", "⣻", "⣽", "⣾"), 

147 clear_with: typing.Optional[str] = None, 

148 _hide_cursor: bool = True, 

149) -> T | ThreadWithReturn[T]: 

150 """ 

151 Provides a pipx style loading animation for a thread. 

152 

153 Args: 

154 thread (ThreadWithReturn): The thread to animate. 

155 text (str): Extra text to show after the spinning icon. 

156 This can be a static value or a callback that's ran at every interval. 

157 threaded (bool): Run the animation in a thread too, unblocking the main thread. 

158 speed (float): The speed of the animation (seconds between animation intervals, defaults to 50ms). 

159 animation (tuple): The frames of the animation. 

160 clear_with (str): replace the animation with a specific character instead of clearing the text line 

161 _hide_cursor (bool): If True, the cursor is hidden during the animation. 

162 

163 Returns: 

164 T: The result of the thread. 

165 """ 

166 with toggle_cursor(enabled=_hide_cursor): 

167 if threaded: 

168 return _animate_threaded(thread, text=text, speed=speed, animation=animation, clear_with=clear_with) 

169 else: 

170 return _animate(thread, text=text, speed=speed, animation=animation, clear_with=clear_with)