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

43 statements  

« prev     ^ index     » next       coverage.py v7.4.3, created at 2024-02-29 21:18 +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 

59@threadify 

60def _animate_threaded( 

61 thread: ThreadWithReturn[T], 

62 speed: float = 0.05, 

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

64) -> T: 

65 return _animate(thread, speed, animation) 

66 

67 

68def _animate( 

69 thread: ThreadWithReturn[T], 

70 speed: float = 0.05, 

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

72) -> T: 

73 """ 

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

75 

76 Args: 

77 thread (ThreadWithReturn): The thread to animate. 

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

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

80 

81 Returns: 

82 T: The result of the thread. 

83 """ 

84 idx = 0 

85 while not thread.is_done(): 

86 idx += 1 

87 print(animation[idx % len(animation)], **_print_kwargs) 

88 time.sleep(speed) 

89 

90 print("\r ", **_print_kwargs) 

91 return thread.result().unwrap() 

92 

93 

94@typing.overload 

95def animate( 

96 thread: ThreadWithReturn[T], 

97 threaded: typing.Literal[True], 

98 speed: float = 0.05, 

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

100 _hide_cursor: bool = True, 

101) -> ThreadWithReturn[T]: 

102 """ 

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

104 """ 

105 

106 

107@typing.overload 

108def animate( 

109 thread: ThreadWithReturn[T], 

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

111 speed: float = 0.05, 

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

113 _hide_cursor: bool = True, 

114) -> T: 

115 """ 

116 Default behavior: run the animation sync. 

117 """ 

118 

119 

120def animate( 

121 thread: ThreadWithReturn[T], 

122 threaded: bool = False, 

123 speed: float = 0.05, 

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

125 _hide_cursor: bool = True, 

126) -> T | ThreadWithReturn[T]: 

127 """ 

128 Provides a pipx style loading animation for a thread. 

129 

130 Args: 

131 thread (ThreadWithReturn): The thread to animate. 

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

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

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

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

136 

137 Returns: 

138 T: The result of the thread. 

139 """ 

140 with toggle_cursor(enabled=_hide_cursor): 

141 if threaded: 

142 return _animate_threaded(thread, speed, animation) 

143 else: 

144 return _animate(thread, speed, animation)