Coverage for src / monte_neo / cli / menu / symbol_selector.py: 0%

1"""Symbol selector with search and multi-column layout.""" 

2 

3from __future__ import annotations 

4 

5import math 

6from typing import TYPE_CHECKING 

7 

8from prompt_toolkit.application import Application 

9from prompt_toolkit.buffer import Buffer 

10from prompt_toolkit.key_binding import KeyBindings 

11from prompt_toolkit.layout.containers import HSplit, ScrollOffsets, Window 

12from prompt_toolkit.layout.controls import FormattedTextControl 

13from prompt_toolkit.layout.layout import Layout 

14from prompt_toolkit.styles import Style as PTStyle 

15from prompt_toolkit.widgets import Frame, TextArea 

16 

17if TYPE_CHECKING: 

18 from prompt_toolkit.key_binding.key_processor import KeyPressEvent 

19 

20 from monte_neo.cli.menu.main import InteractiveMenu 

21 

22 

23class SymbolSelector: 

24 """A searchable grid-based symbol selector.""" 

25 

26 def __init__(self, symbols: list[str], style=None): 

27 self.symbols = symbols 

28 self.filtered_symbols = symbols 

29 # Convert questionary style to prompt_toolkit style if needed 

30 if style and hasattr(style, "class_names_and_attrs"): 

31 self.pt_style = style 

32 elif isinstance(style, list): 

33 self.pt_style = PTStyle(style) 

34 else: 

35 self.pt_style = None 

36 

37 self.index = 0 

38 self.cols = 4 # Initial value, will be updated based on width 

39 self.col_width = 18 # Symbol width + padding 

40 self.search_text = "" 

41 self.result: str | None = None 

42 

43 # UI components 

44 self.search_field = TextArea( 

45 prompt="Search: ", 

46 multiline=False, 

47 ) 

48 # Add handler to buffer 

49 self.search_field.buffer.on_text_changed += self._on_search_change 

50 

51 self.grid_control = FormattedTextControl( 

52 text=self._get_grid_text, 

53 focusable=True, 

54 # This is crucial for scrolling: we tell prompt_toolkit where the "cursor" is 

55 get_cursor_position=self._get_cursor_position, 

56 ) 

57 self.grid_window = Window( 

58 self.grid_control, 

59 height=15, 

60 # Enable scrolling when cursor goes out of view 

61 scroll_offsets=ScrollOffsets(top=1, bottom=1), 

62 ) 

63 

64 self.kb = KeyBindings() 

65 self._setup_keybindings() 

66 

67 self.app = Application( 

68 layout=Layout( 

69 HSplit([ 

70 Frame(self.search_field, title="Type to Search"), 

71 Frame(self.grid_window, title=self._get_title), 

72 ]), 

73 focused_element=self.search_field, 

74 ), 

75 key_bindings=self.kb, 

76 style=self.pt_style, 

77 full_screen=False, 

78 mouse_support=True, 

79 ) 

80 

81 def _get_title(self) -> str: 

82 """Return dynamic title with counts.""" 

83 total = len(self.symbols) 

84 filtered = len(self.filtered_symbols) 

85 return f"Select Symbol ({filtered}/{total} tokens) (Arrows to navigate, Enter to select)" 

86 

87 def _get_cursor_position(self): 

88 """Return the (x, y) position of the currently selected symbol for scrolling.""" 

89 if not self.filtered_symbols: 

90 return None 

91 

92 # Calculate row and column based on current index 

93 row = self.index // self.cols 

94 col = self.index % self.cols 

95 

96 # x is the character position in the line 

97 x = col * self.col_width 

98 # y is the line number 

99 y = row 

100 

101 from prompt_toolkit.data_structures import Point 

102 return Point(x=x, y=y) 

103 

104 def _update_cols(self): 

105 """Update columns count based on available window width.""" 

106 width = self.app.renderer.output.get_size().columns 

107 # Subtract some padding for frames and margins 

108 available_width = max(20, width - 10) 

109 self.cols = max(1, available_width // self.col_width) 

110 

111 def _on_search_change(self, buffer: Buffer) -> None: 

112 self.search_text = buffer.text.upper() 

113 self.filtered_symbols = [s for s in self.symbols if self.search_text in s] 

114 # Adjust index if out of bounds after filtering 

115 if not self.filtered_symbols: 

116 self.index = 0 

117 else: 

118 self.index = min(self.index, len(self.filtered_symbols) - 1) 

119 

120 def _get_grid_text(self): 

121 if not self.filtered_symbols: 

122 return [("class:disabled", " No matches found")] 

123 

124 # Update columns based on current window size 

125 self._update_cols() 

126 

127 rows = math.ceil(len(self.filtered_symbols) / self.cols) 

128 result = [] 

129 for r in range(rows): 

130 line = [] 

131 for c in range(self.cols): 

132 idx = r * self.cols + c 

133 if idx < len(self.filtered_symbols): 

134 symbol = self.filtered_symbols[idx] 

135 # Ensure each column has fixed width for cursor positioning 

136 text = f" {symbol:<15} " 

137 if idx == self.index: 

138 line.append(("class:highlighted", text)) 

139 else: 

140 line.append(("class:text", text)) 

141 result.extend(line) 

142 result.append(("", "\n")) 

143 return result 

144 

145 def _setup_keybindings(self): 

146 @self.kb.add("up") 

147 def _(event: KeyPressEvent): 

148 if self.index >= self.cols: 

149 self.index -= self.cols 

150 

151 @self.kb.add("down") 

152 def _(event: KeyPressEvent): 

153 if self.index + self.cols < len(self.filtered_symbols): 

154 self.index += self.cols 

155 

156 @self.kb.add("left") 

157 def _(event: KeyPressEvent): 

158 if self.index > 0: 

159 self.index -= 1 

160 

161 @self.kb.add("right") 

162 def _(event: KeyPressEvent): 

163 if self.index < len(self.filtered_symbols) - 1: 

164 self.index += 1 

165 

166 @self.kb.add("enter") 

167 def _(event: KeyPressEvent): 

168 if self.filtered_symbols: 

169 self.result = self.filtered_symbols[self.index] 

170 event.app.exit(result=self.result) 

171 

172 @self.kb.add("escape") 

173 @self.kb.add("c-c") 

174 def _(event: KeyPressEvent): 

175 self.result = None 

176 event.app.exit() 

177 

178 def ask(self) -> str | None: 

179 """Run the selector and return the selected symbol.""" 

180 self.app.run() 

181 return self.result 

182 

183def select_symbol(menu: InteractiveMenu) -> str | None: 

184 """Helper function to run the symbol selector.""" 

185 import os 

186 

187 # Get available symbols from data directory 

188 data_dir = menu.config.data_dir 

189 raw_dir = os.path.join(data_dir, "raw") 

190 if not os.path.exists(raw_dir): 

191 return None 

192 

193 symbols = [] 

194 for f in os.listdir(raw_dir): 

195 if f.endswith(".parquet"): 

196 # Format: SYMBOL_TIMEFRAME.parquet 

197 symbol_part = f.split("_")[0] 

198 if symbol_part not in symbols: 

199 symbols.append(symbol_part) 

200 

201 if not symbols: 

202 return None 

203 

204 from monte_neo.cli.styles import CUSTOM_STYLE 

205 selector = SymbolSelector(symbols, style=CUSTOM_STYLE) 

206 return selector.ask()