Coverage for /builds/ase/ase/ase/gui/atomseditor.py: 92.14%

1from dataclasses import dataclass 

2from typing import Callable 

3 

4import numpy as np 

5 

6import ase.gui.ui as ui 

7from ase.gui.i18n import _ 

8 

9 

10@dataclass 

11class Column: 

12 name: str 

13 displayname: str 

14 widget_width: int 

15 getvalue: Callable 

16 setvalue: Callable 

17 format_value: Callable = lambda obj: str(obj) 

18 

19 

20class AtomsEditor: 

21 # We subscribe to gui.draw() calls in order to track changes, 

22 # but we should have an actual "atoms changed" event instead. 

23 

24 def __init__(self, gui): 

25 gui.obs.change_atoms.register(self.update_table_from_atoms) 

26 

27 win = ui.Window(_('Edit atoms'), wmtype='utility') 

28 

29 treeview = ui.ttk.Treeview(win.win, selectmode='extended') 

30 edit_entry = ui.ttk.Entry(win.win) 

31 edit_entry.pack(side='bottom', fill='x') 

32 treeview.pack(side='left', fill='y') 

33 bar = ui.ttk.Scrollbar( 

34 win.win, orient='vertical', command=self.scroll_via_scrollbar 

35 ) 

36 treeview.configure(yscrollcommand=self.scroll_via_treeview) 

37 

38 treeview.column('#0', width=40) 

39 treeview.heading('#0', text=_('id')) 

40 

41 bar.pack(side='right', fill='y') 

42 self.scrollbar = bar 

43 

44 def get_symbol(atoms, i): 

45 return atoms.symbols[i] 

46 

47 def set_symbol(atoms, i, value): 

48 from ase.data import atomic_numbers 

49 

50 if value not in atomic_numbers: 

51 return # Display error? 

52 atoms.symbols[i] = value 

53 

54 self.gui = gui 

55 self.treeview = treeview 

56 self._current_entry = None 

57 

58 columns = [] 

59 symbols_column = Column( 

60 'symbol', _('symbol'), 60, get_symbol, set_symbol 

61 ) 

62 columns.append(symbols_column) 

63 

64 class GetSetPos: 

65 def __init__(self, c): 

66 self.c = c 

67 

68 def set_position(self, atoms, i, value): 

69 try: 

70 value = float(value) 

71 except ValueError: 

72 return 

73 atoms.positions[i, self.c] = value 

74 

75 def get_position(self, atoms, i): 

76 return atoms.positions[i, self.c] 

77 

78 for c, axisname in enumerate('xyz'): 

79 column = Column( 

80 axisname, 

81 axisname, 

82 92, 

83 GetSetPos(c).get_position, 

84 GetSetPos(c).set_position, 

85 format_value=lambda val: f'{val:.4f}', 

86 ) 

87 columns.append(column) 

88 

89 self.columns = columns 

90 

91 treeview.bind('<Double-1>', self.doubleclick) 

92 treeview.bind('<<TreeviewSelect>>', self.treeview_selection_changed) 

93 

94 self.define_columns_on_widget() 

95 self.update_table_from_atoms() 

96 

97 self.edit_entry = edit_entry 

98 

99 def treeview_selection_changed(self, event): 

100 selected_items = self.treeview.selection() 

101 indices = [self.rownumber(item) for item in selected_items] 

102 self.gui.set_selected_atoms(indices) 

103 

104 def scroll_via_scrollbar(self, *args, **kwargs): 

105 self.leave_edit_mode() 

106 return self.treeview.yview(*args, **kwargs) 

107 

108 def scroll_via_treeview(self, *args, **kwargs): 

109 # Here it is important to leave edit mode since scrolling 

110 # invalidates the widget location. Alternatively we could keep 

111 # it open as long as we move it but that sounds like work 

112 self.leave_edit_mode() 

113 return self.scrollbar.set(*args, **kwargs) 

114 

115 def leave_edit_mode(self): 

116 if self._current_entry is not None: 

117 self._current_entry.destroy() 

118 self._current_entry = None 

119 self.treeview.focus_force() 

120 

121 @property 

122 def atoms(self): 

123 return self.gui.atoms 

124 

125 def update_table_from_atoms(self): 

126 self.treeview.delete(*self.treeview.get_children()) 

127 for i in range(len(self.atoms)): 

128 values = self.get_row_values(i) 

129 self.treeview.insert( 

130 '', 'end', text=i, values=values, iid=self.rowid(i) 

131 ) 

132 

133 mask = self.gui.images.selected[: len(self.atoms)] 

134 selection = np.arange(len(self.atoms))[mask] 

135 

136 rowids = [self.rowid(index) for index in selection] 

137 # Note: selection_set() does *not* fire an event, and therefore 

138 # we do not need to worry about infinite recursion. 

139 # However the event listening is wonky now because we need 

140 # better GUI change listeners. 

141 self.treeview.selection_set(*rowids) 

142 

143 def get_row_values(self, i): 

144 return [ 

145 column.format_value(column.getvalue(self.atoms, i)) 

146 for column in self.columns 

147 ] 

148 

149 def define_columns_on_widget(self): 

150 self.treeview['columns'] = [column.name for column in self.columns] 

151 for column in self.columns: 

152 self.treeview.heading(column.name, text=column.displayname) 

153 self.treeview.column( 

154 column.name, 

155 width=column.widget_width, 

156 anchor='e', 

157 ) 

158 

159 def rowid(self, rownumber: int) -> str: 

160 return f'R{rownumber}' 

161 

162 def rownumber(self, rowid: str) -> int: 

163 assert rowid.startswith('R'), repr(rowid) 

164 return int(rowid[1:]) 

165 

166 def set_value(self, column_no: int, row_no: int, value: object) -> None: 

167 column = self.columns[column_no] 

168 column.setvalue(self.atoms, row_no, value) 

169 text = column.format_value(column.getvalue(self.atoms, row_no)) 

170 

171 # The text that we set here is not what matters: It may be rounded. 

172 # It was column.setvalue() which did the actual change. 

173 self.treeview.set(self.rowid(row_no), column.name, value=text) 

174 

175 # (Maybe it is not always necessary to redraw everything.) 

176 self.gui.set_frame() 

177 

178 def doubleclick(self, event): 

179 row_id = self.treeview.identify_row(event.y) 

180 column_id = self.treeview.identify_column(event.x) 

181 if not row_id or not column_id: 

182 return # clicked outside actual rows/columns 

183 self.edit_field(row_id, column_id) 

184 

185 def edit_field(self, row_id, column_id): 

186 assert column_id.startswith('#'), repr(column_id) 

187 column_no = int(column_id[1:]) - 1 

188 

189 if column_no == -1: 

190 return # This is the ID column. 

191 

192 row_no = self.rownumber(row_id) 

193 assert 0 <= column_no < len(self.columns) 

194 assert 0 <= row_no < len(self.atoms) 

195 

196 content = self.columns[column_no].getvalue(self.atoms, row_no) 

197 

198 assert self._current_entry is None 

199 entry = ui.ttk.Entry(self.treeview) 

200 entry.insert(0, content) 

201 entry.focus_force() 

202 entry.selection_range(0, 'end') 

203 

204 def apply_change(_event=None): 

205 value = entry.get() 

206 try: 

207 self.set_value(column_no, row_no, value) 

208 finally: 

209 # Focus was given to the text field, now return it: 

210 self.treeview.focus_force() 

211 self.leave_edit_mode() 

212 

213 entry.bind('<FocusOut>', apply_change) 

214 ui.bind_enter(entry, apply_change) 

215 entry.bind('<Escape>', lambda *args: self.leave_edit_mode()) 

216 

217 bbox = self.treeview.bbox(row_id, column_id) 

218 if bbox: # (bbox is '' when testing without display) 

219 x, y, width, height = bbox 

220 entry.place(x=x, y=y, height=height) 

221 self._current_entry = entry 

222 return entry, apply_change