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