1import os 2 3from pymol.Qt import QtGui, QtCore, QtWidgets 4from pymol.Qt.utils import UpdateLock, PopupOnException 5import pymol 6 7class UneditableDelegate(QtWidgets.QStyledItemDelegate): 8 def createEditor(self, parent, option, index): 9 return None 10 11class FunctionSuspender: 12 def __init__(self, func): 13 self.func = func 14 def __enter__(self): 15 self.func.suspended = True 16 def __exit__(self, exc_type, exc_value, traceback): 17 self.func.suspended = False 18 19def suspendable(func): 20 def wrapper(*args, **kwargs): 21 if not func.suspended: 22 return func(*args, **kwargs) 23 func.suspended = False 24 wrapper.suspend = FunctionSuspender(func) 25 return wrapper 26 27def get_object_names(_self): 28 # was _self.get_names('public_objects') in PyMOL 2.1 29 # but that throws exceptions for groups/isosurfaces/etc. 30 names = _self.get_object_list() 31 return names 32 33def props_dialog(parent): #noqa 34 from pymol.setting import name_dict 35 36 cmd = parent.cmd 37 form = parent.load_form('props', 'floating') 38 parent.addDockWidget(QtCore.Qt.RightDockWidgetArea, form._dialog) 39 40 def make_entry(parent, label): 41 item = QtWidgets.QTreeWidgetItem(parent) 42 item.setText(0, str(label)) 43 item.setFlags(QtCore.Qt.ItemIsEditable | 44 QtCore.Qt.ItemIsEnabled | 45 QtCore.Qt.ItemIsSelectable ) 46 return item 47 48 def make_cat(parent, label): 49 item = QtWidgets.QTreeWidgetItem(parent) 50 item.setText(0, str(label)) 51 item.setFirstColumnSpanned(True) 52 item.setExpanded(True) 53 item.setChildIndicatorPolicy( 54 QtWidgets.QTreeWidgetItem.ShowIndicator) 55 return item 56 57 # make first column uneditable 58 form.treeWidget.setItemDelegateForColumn(0, UneditableDelegate()) 59 60 item_object = make_cat(form.treeWidget, "Object-Level") 61 item_object_ttt = make_entry(item_object, "TTT Matrix") 62 item_object_settings = make_cat(item_object, "Settings") 63 64 item_ostate = make_cat(form.treeWidget, "Object-State-Level") 65 item_ostate_title = make_entry(item_ostate, "Title") 66 item_ostate_matrix = make_entry(item_ostate, "State Matrix") 67 item_ostate_settings = make_cat(item_ostate, "Settings") 68 item_ostate_properties = Ellipsis # Incentive PyMOL only 69 70 item_atom = make_cat(form.treeWidget, "Atom-Level") 71 item_atom_identifiers = make_cat(item_atom, "Identifiers") 72 item_atom_builtins = make_cat(item_atom, "Properties (built-in)") 73 item_atom_settings = make_cat(item_atom, "Settings") 74 item_atom_properties = Ellipsis # Incentive PyMOL only 75 76 item_astate = make_cat(form.treeWidget, "Atom-State-Level") 77 item_astate_builtins = make_cat(item_astate, "Properties (built-in)") 78 item_astate_settings = make_cat(item_astate, "Settings") 79 80 keys_atom_identifiers = ['model', 'index', 'segi', 'chain', 'resi', 81 'resn', 'name', 'alt', 'ID', 'rank'] 82 keys_atom_builtins = ['elem', 'q', 'b', 'type', 'formal_charge', 83 'partial_charge', 'numeric_type', 'text_type', 84 # avoid stereo auto-assignment errors 85 # 'stereo', 86 'vdw', 'ss', 'color', 'reps', 87 'protons', 'geom', 'valence', 'elec_radius'] 88 keys_astate_builtins = ['state', 'x', 'y', 'z'] 89 90 items = {} 91 for key in keys_atom_identifiers: 92 items[key] = make_entry(item_atom_identifiers, key) 93 for key in keys_atom_builtins: 94 items[key] = make_entry(item_atom_builtins, key) 95 for key in keys_astate_builtins: 96 items[key] = make_entry(item_astate_builtins, key) 97 98 items['model'].setDisabled(True) 99 items['index'].setDisabled(True) 100 items['state'].setDisabled(True) 101 102 def item_changed(item, column): 103 """ 104 Edits current item. 105 """ 106 107 if item_changed.skip: 108 return 109 110 model = form.input_model.currentText() 111 state = form.input_state.value() 112 key = item.text(0) 113 new_value = item.text(1) 114 parent = item.parent() 115 116 result = False 117 if item is item_object_ttt: 118 try: 119 if new_value: 120 result = cmd.set_object_ttt(model, new_value) 121 except (ValueError, IndexError): 122 result = False 123 elif item is item_ostate_title: 124 result = cmd.set_title(model, state, new_value) 125 elif item is item_ostate_matrix: 126 cmd.matrix_reset(model, state) 127 try: 128 new_value = cmd.safe_eval(new_value) 129 result = cmd.transform_object(model, new_value, state) 130 except: # CmdTransformObject-DEBUG: bad matrix 131 result = False 132 elif parent is item_object_settings: 133 with PopupOnException(): 134 cmd.set(key, new_value, model, quiet=0) 135 elif parent is item_ostate_settings: 136 with PopupOnException(): 137 cmd.set(key, new_value, model, state, quiet=0) 138 elif parent is item_ostate_properties: 139 cmd.set_property(key, new_value, model, state, quiet=0) 140 else: 141 is_state = False 142 143 if parent is item_atom_properties: 144 key = 'p.' + key 145 elif parent is item_atom_settings: 146 key = 's.' + key 147 elif parent is item_astate_settings: 148 key = 's.' + key 149 is_state = True 150 elif key in keys_astate_builtins: 151 is_state = True 152 153 def get_new_value(old_value): 154 if isinstance(old_value, (tuple, list)): 155 return cmd.safe_eval(new_value) 156 157 try: 158 # cast to old type (required for e.g. 'resv = "3"') 159 return type(old_value)(new_value) 160 except ValueError: 161 # return str and let PyMOL handle it (e.g. change 162 # type of user property) 163 return new_value.encode('ascii') 164 165 alter_args = ('pk1', key + '= get_new_value(' + key + ')', 0, 166 {'get_new_value': get_new_value}) 167 168 if is_state: 169 result = cmd.alter_state(state, *alter_args) 170 else: 171 result = cmd.alter(*alter_args) 172 173 if not result: 174 update_treewidget_model() 175 176 item_changed.skip = False 177 178 def update_object_settings(parent, model, state): 179 parent.takeChildren() 180 for sitem in (cmd.get_object_settings(model, state) or []): 181 key = name_dict.get(sitem[0], sitem[0]) 182 item = make_entry(parent, key) 183 item.setText(1, str(sitem[2])) 184 185 def update_atom_settings(wrapper, parent): 186 parent.takeChildren() 187 for key in wrapper: 188 item = make_entry(parent, name_dict.get(key, key)) 189 value = wrapper[key] 190 item.setText(1, str(value)) 191 192 def update_atom_properties(wrapper, parent): 193 parent.takeChildren() 194 for key in wrapper: 195 item = make_entry(parent, key) 196 value = wrapper[key] 197 item.setText(1, str(value)) 198 199 def update_atom_fields(ns): 200 for key in keys_atom_identifiers + keys_atom_builtins: 201 try: 202 value = ns[key] 203 except Exception as e: 204 value = 'ERROR: ' + str(e) 205 items[key].setText(1, str(value)) 206 update_atom_settings(ns['s'], item_atom_settings) 207 208 def update_astate_fields(ns): 209 for key in keys_astate_builtins: 210 value = ns[key] 211 items[key].setText(1, str(value)) 212 update_atom_settings(ns['s'], item_astate_settings) 213 214 space = { 215 'update_atom_fields': update_atom_fields, 216 'update_astate_fields': update_astate_fields, 217 'locals': locals, 218 } 219 220 def update_from_pk1(): 221 pk1_atom = [] 222 if cmd.iterate('?pk1', 'pk1_atom[:] = [model, index]', space=locals()) > 0: 223 form.input_model.setCurrentIndex(form.input_model.findText(pk1_atom[0])) 224 form.input_index.setValue(pk1_atom[1]) 225 226 def update_pk1(): 227 model = form.input_model.currentText() 228 index = form.input_index.value() 229 230 if model and index: 231 try: 232 cmd.edit((model, index)) 233 return True 234 except pymol.CmdException: 235 pass 236 237 return False 238 239 def update_treewidget(*args): 240 if not update_pk1(): 241 return 242 243 state = form.input_state.value() 244 245 item_changed.skip = True 246 count = cmd.iterate( 247 'pk1', 'update_atom_fields(locals())', space=space) 248 item_atom.setDisabled(not count) 249 if count: 250 count = cmd.iterate_state( 251 state, 'pk1', 252 'update_astate_fields(locals())', space=space) 253 item_astate.setDisabled(not count) 254 item_changed.skip = False 255 256 def update_treewidget_state(*args): 257 model = form.input_model.currentText() 258 state = form.input_state.value() 259 if not (model and state): 260 return 261 262 item_changed.skip = True 263 item_ostate_title.setText( 264 1, str(cmd.get_title(model, state) or '')) 265 item_ostate_matrix.setText( 266 1, str( 267 cmd.get_object_matrix( 268 model, state, 0) or '')) 269 270 update_object_settings(item_ostate_settings, model, state) 271 272 update_treewidget() 273 item_changed.skip = False 274 275 @suspendable 276 def update_treewidget_model(*args): 277 item_changed.skip = True 278 model = form.input_model.currentText() 279 280 if not model: 281 return 282 283 item_object_ttt.setText(1, str(cmd.get_object_ttt(model) or '')) 284 update_object_settings(item_object_settings, model, 0) 285 286 natoms = cmd.count_atoms('?' + model) 287 nstates = cmd.count_states('?' + model) 288 289 form.input_state.setMinimum(1) 290 form.input_state.setMaximum(nstates) 291 form.input_index.setMinimum(1) 292 form.input_index.setMaximum(natoms) 293 294 item_atom.setHidden(natoms == 0) 295 item_astate.setHidden(natoms == 0) 296 item_ostate.setHidden(nstates == 0) 297 298 update_treewidget_state() 299 item_changed.skip = False 300 301 def update_model_list(*args): 302 if 'pk1' not in cmd.get_names('selections'): 303 update_pk1() 304 305 with update_treewidget_model.suspend: 306 form.input_model.clear() 307 form.input_model.addItems(get_object_names(cmd)) 308 update_from_pk1() 309 310 update_treewidget_model() 311 312 # init input fields 313 form.input_model.addItems(get_object_names(cmd)) 314 form.input_state.setValue(cmd.get_state()) 315 316 # select pk1 atom if available 317 update_from_pk1() 318 319 # hook up events 320 form.input_model.currentIndexChanged.connect(update_treewidget_model) 321 form.input_state.valueChanged.connect(update_treewidget_state) 322 form.input_index.valueChanged.connect(update_treewidget) 323 form.button_refresh.clicked.connect(update_model_list) 324 325 # themed icons only available by default on X11 326 if form.button_refresh.icon().isNull(): 327 form.button_refresh.setIcon(QtGui.QIcon( 328 os.path.expandvars('$PYMOL_DATA/pmg_qt/icons/refresh.svg'))) 329 330 # update and show 331 update_treewidget_model() 332 form.treeWidget.setColumnWidth(0, 200) 333 334 form.treeWidget.itemChanged.connect(item_changed) 335 336 return form._dialog 337