import os
import textwrap
import schrodinger
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.structutils import analyze
from schrodinger.trajectory.trajectory_gui_dir.export_structure_enums import \
ExportFrameOption
from schrodinger.trajectory.trajectory_gui_dir.export_structure_enums import \
ExportToOption
from schrodinger.ui import maestro_ui
from schrodinger.ui.qt import atomselector
from schrodinger.ui.qt import filedialog
from schrodinger.ui.qt import utils
from schrodinger.utils import fileutils
from . import export_structures_ui
from . import stylesheet
maestro = schrodinger.get_maestro()
[docs]class ExportStructures(QtWidgets.QDialog):
"""
Export structures class for trajectory viewer
:cvar exportButtonClicked: A signal emitted when clicked on 'Export' button.
:vartype exportButtonClicked: `QtCore.pyqtSignal`
"""
exportButtonClicked = QtCore.pyqtSignal()
[docs] def __init__(self, parent: QtWidgets.QWidget):
"""
:param parent: Parent widget.
"""
super(QtWidgets.QDialog, self).__init__(parent)
self.ui = export_structures_ui.Ui_Dialog()
self.setStyleSheet(stylesheet.EXPORT_STRUCTURES_DIALOG_STYLESHEET)
self.ui.setupUi(self)
# Construct Atom Selector and place it in the place holder
self.atom_selector = atomselector.AtomSelector(self,
show_pick=False,
show_selection=False,
show_plus=True)
self.atom_selector.main_layout.setContentsMargins(2, 2, 2, 2)
layout = self.ui.atomselector_widget.layout()
layout.insertWidget(0, self.atom_selector)
self.atom_selector.atomSelectionDialogDismissed.connect(self.show)
self.atom_selector.atomSelectionDialogAboutToBeShown.connect(self.hide)
self.adjustSize()
self.ui.pt_radiobutton.setChecked(True)
self.ui.file_lineedit.setEnabled(False)
self.ui.browse_pushbutton.setEnabled(False)
self.ui.export_buttongroup.buttonToggled.connect(
self.updateFileDependentOptions)
self.ui.browse_pushbutton.clicked.connect(self.browse)
self.ui.one_radiobutton.toggled.connect(
self.ui.one_entries_label.setEnabled)
self.ui.all_range_radiobutton.toggled.connect(
self.ui.all_entries_label.setEnabled)
self.ui.current_radiobutton.toggled.connect(
self.ui.current_label.setEnabled)
self.ui.one_radiobutton.setChecked(True)
self.ui.all_entries_label.setEnabled(False)
self.ui.current_label.setEnabled(False)
self.ui.frames_buttongroup.buttonToggled.connect(
self.updateFramesStackedWidget)
self.ui.split_by_checkbox.toggled.connect(self.splitByComponentChanged)
self.ui.limit_checkbox.toggled.connect(self.limitRangeToggled)
self.ui.reset_button.clicked.connect(self.resetRange)
self.ui.start_lineedit.setValidator(QtGui.QIntValidator(self))
self.ui.end_lineedit.setValidator(QtGui.QIntValidator(self))
self.ui.start_lineedit.editingFinished.connect(self.startChanged)
self.ui.start_lineedit.textEdited.connect(self.startChanged)
self.ui.end_lineedit.editingFinished.connect(self.endChanged)
self.ui.end_lineedit.textEdited.connect(self.endChanged)
self.ui.all_radiobutton.setChecked(True)
self.ui.atomselector_widget.setEnabled(False)
self.ui.specified_radiobutton.toggled.connect(
self.ui.atomselector_widget.setEnabled)
self.ui.load_sel_pushbutton.clicked.connect(
self.atom_selector.loadWorkspaceSelection)
self.ui.export_buttongroup.buttonToggled.connect(
self.updateMultiStructureOptions)
self.ui.export_buttongroup.buttonToggled.connect(
self.updateExportButton)
self.ui.structures_buttongroup.buttonToggled.connect(
self.updateExportButton)
self.ui.file_lineedit.textChanged.connect(
self.updateMultiStructureOptions)
self.ui.file_lineedit.textChanged.connect(self.updateExportButton)
self.atom_selector.aslTextModified.connect(self.updateExportButton)
self.ui.cancel_pushbutton.clicked.connect(self.cancelClicked)
self.ui.export_pushbutton.clicked.connect(self.exportClicked)
self.ui.help_button.clicked.connect(self.help)
self._last_dir = None
self.file_formats = [fileutils.MAESTRO_STRICT, fileutils.CMS]
self.ui.info_button.setToolTip(
textwrap.dedent("""
Use the Maestro format to export multiple
frames or separate components of the
current frame. Other formats are available
for a single entry only.
"""))
[docs] def updateFileDependentOptions(self):
"""
Update File radio option dependent options.
"""
checked = self.ui.file_radiobutton.isChecked()
self.ui.file_lineedit.setEnabled(checked)
self.ui.browse_pushbutton.setEnabled(checked)
self.ui.info_button.setVisible(checked)
[docs] def updateMultiStructureOptions(self):
"""
Update options in the dialog.
If user selects cms file, we should not allow user to export
more than one structure, so disable all components related to multiple
structures.
"""
file_path = self.ui.file_lineedit.text()
only_single_st = (self.ui.file_radiobutton.isChecked() and
fileutils.is_cms_file(file_path))
self.ui.one_radiobutton.setEnabled(not only_single_st and
self._step_size > 1)
self.ui.all_range_radiobutton.setEnabled(not only_single_st)
self.ui.split_by_checkbox.setEnabled(not only_single_st)
if only_single_st:
self.ui.current_radiobutton.setChecked(True)
self.ui.split_by_checkbox.setChecked(False)
self.ui.specified_radiobutton.setEnabled(False)
else:
self.updateSpecifiedASLButton()
[docs] def browse(self):
"""
Slot triggered when 'Browse...' button is clicked on
"""
filter_string = filedialog.filter_string_from_formats(self.file_formats)
file_name = filedialog.get_save_file_name(
self,
dir=self._last_dir,
caption='Choose File for Export',
filter=filter_string)
if file_name:
self._last_dir = os.path.dirname(file_name)
# Ok pressed
if file_name:
self.ui.file_lineedit.setText(str(file_name))
[docs] def splitByComponentChanged(self, checked):
"""
Slot triggered when 'Split by component' is toggled
"""
total_entries = 1 + (self._comp_cts_total if checked else 0)
verb = 'entry' if total_entries == 1 else 'entries'
label = f'({total_entries} {verb})'
self.ui.current_label.setText(label)
# If split by component option is turned on, always disable specified
# atoms asl radio button and enable all radio button.
self.updateSpecifiedASLButton()
if checked:
self.ui.all_radiobutton.setChecked(True)
[docs] def showDlg(self, start_frame, end_frame, step_size, total_frame,
comp_cts_total):
"""
Show 'Export Structures' dialog with given information
:param start_frame: Start frame
:type start_frame: int
:param end_frame: End frame
:type end_frame: int
:param step_size: Step size
:type step_size: int
:param total_frame: Total frames
:type total_frame: int
:param comp_cts_total: Number of component cts.
:type comp_cts_total: int
"""
ws_hub = maestro_ui.WorkspaceHub.instance()
has_ws_atoms = ws_hub.getSelectedAtomCount() > 0
self.initDlg(start_frame, end_frame, step_size, total_frame,
comp_cts_total, has_ws_atoms)
# Show the dialog
self.show()
[docs] def initDlg(self, start_frame, end_frame, step_size, total_frame,
comp_cts_total, has_ws_atoms):
"""
Initialze 'Export Structures' dialog with given information
:param start_frame: Start frame
:type start_frame: int
:param end_frame: End frame
:type end_frame: int
:param step_size: Step size
:type step_size: int
:param total_frame: Total frames
:type total_frame: int
:param comp_cts_total: Number of component cts.
:type comp_cts_total: int
:param has_ws_atoms: True if workspace has atoms.
:type has_ws_atoms: bool
"""
# Enable/disable 'One per step' option based on step size
self.ui.one_radiobutton.setEnabled(step_size > 1)
# Change the default if 'One per step' is selected and is disabled
if step_size == 1 and self.ui.one_radiobutton.isChecked():
self.ui.all_range_radiobutton.setChecked(True)
self.start = self._start_frame = start_frame
self.end = self._end_frame = end_frame
self._total_frame = total_frame
self._step_size = step_size
self._comp_cts_total = comp_cts_total
self.ui.start_lineedit.validator().setRange(1, self._total_frame - 1)
self.ui.end_lineedit.validator().setRange(2, self._total_frame)
self.splitByComponentChanged(self.ui.split_by_checkbox.isChecked())
# Reset range according to given values
self.resetRange()
# Enable 'Load Selection' only if workspace has selection
self.ui.load_sel_pushbutton.setEnabled(has_ws_atoms)
# Save current option values
self.saveOptions()
# If we are showing dialog second time, we might have file extension as
# .cms, so we should update multi structure related options.
self.updateMultiStructureOptions()
self.updateFileDependentOptions()
# Update 'Export' button
self.updateExportButton()
[docs] def saveOptions(self):
"""
Save current option values
"""
self._export_button = self.ui.export_buttongroup.checkedButton()
self._export_file_name = self.ui.file_lineedit.text()
self._frames_button = self.ui.frames_buttongroup.checkedButton()
self._split_by_component = self.ui.split_by_checkbox.isChecked()
self._structures_button = self.ui.structures_buttongroup.checkedButton()
self._structures_asl = self.atom_selector.getAsl()
[docs] def isValidASL(self):
"""
Return True if asl in the text box is valid.
"""
return analyze.validate_asl(self.export_asl)
[docs] def exportClicked(self):
"""
Accept and emit exportButtonClicked signal when 'Export' button is
clicked
"""
if not self.isValidASL():
asl = self.export_asl
maestro.warning(f"Invalid ASL {asl}")
return
self.close()
self.exportButtonClicked.emit()
[docs] def cancelClicked(self):
"""
Restore saved option values
"""
self.close()
self._export_button.setChecked(True)
self.ui.file_lineedit.setText(self._export_file_name)
self._frames_button.setChecked(True)
self.ui.split_by_checkbox.setChecked(self._split_by_component)
self._structures_button.setChecked(True)
self.atom_selector.setAsl(self._structures_asl)
def _getIntFromText(self, value):
"""
Return integer value of given string value.
:param value: Value to be converted
:type value: str
:rtype: int
"""
try:
val = int(value)
except ValueError:
val = 0
return val
@property
def export_frame_option(self):
"""
:return: Return frame option
:rtype: enum(ExportFrameOption)
"""
if self.ui.one_radiobutton.isChecked():
return ExportFrameOption.ONE_PER_STEP_IN_RANGE
elif self.ui.all_range_radiobutton.isChecked():
if self.ui.limit_checkbox.isChecked():
return ExportFrameOption.ALL_IN_LIMITED_RANGE
else:
return ExportFrameOption.ALL_IN_RANGE
elif self.ui.split_by_checkbox.isChecked():
return ExportFrameOption.CURRENT_FRAME_WITH_COMPONENT_CTS
else:
return ExportFrameOption.ONLY_CURRENT_FRAME
@property
def export_to_option(self):
"""
:return: Return export option
:rtype: enum(ExportOption)
"""
if self.ui.pt_radiobutton.isChecked():
return ExportToOption.PROJECT_TABLE
else:
return ExportToOption.EXTERNAL_FILE
@property
def export_file_path(self):
"""
If file path does not contain extension, it also adds default .mae
extension.
:rtype: str
:return: Return file path to be used for export.
"""
file_path = self.ui.file_lineedit.text()
if (fileutils.is_cms_file(file_path) or
fileutils.is_maestro_file(file_path)):
return file_path
else:
return file_path + ".mae"
@property
def export_asl(self):
"""
:rtype: str
:return: Return asl of atoms to be exported.
"""
if self.ui.specified_radiobutton.isChecked():
return self.atom_selector.getAsl()
else:
return "all"
@property
def start(self):
"""
Return 'Start' value
"""
return self._getIntFromText(self.ui.start_lineedit.text())
@start.setter
def start(self, value):
"""
Set 'End' value
"""
self.ui.start_lineedit.setText(str(value))
@property
def end(self):
"""
Return 'End' value
"""
return self._getIntFromText(self.ui.end_lineedit.text())
@end.setter
def end(self, value):
"""
Set 'End' value
"""
self.ui.end_lineedit.setText(str(value))
@property
def frame_count(self):
"""
Return frame count
"""
return self.end - self.start + 1
[docs] def updateEntryCountLabels(self):
"""
Updates entry count labels according to current values
"""
s = lambda n: 'entry' if n == 1 else 'entries'
count_by_step = len(range(self.start, self.end + 1, self._step_size))
self.ui.one_entries_label.setText('(%d %s)' %
(count_by_step, s(count_by_step)))
self.ui.all_entries_label.setText(
'(%d %s)' % (self.frame_count, s(self.frame_count)))
[docs] def isValidStart(self, start=None, end=None):
"""
Whether start is valid
:param start: If None then current start will be taken.
:type value: int
:param end: If None then current end will be taken.
:type value: int
:rtype: bool
"""
start = start if start else self.start
end = end if end else self.end
return 0 < start < end
[docs] def isValidEnd(self, start=None, end=None):
"""
Whether end is valid
:param start: If None then current start will be taken.
:type value: int
:param end: If None then current end will be taken.
:type value: int
:rtype: bool
"""
start = start if start else self.start
end = end if end else self.end
return 0 < end > start
[docs] def limitRangeToggled(self, checked):
"""
Triggered when 'Limit range' toggled
"""
# Validate self._last_start and reset if required
if not self.isValidStart(self._last_start, self._last_end):
self._last_start = self._start_frame
# Validate self._last_end and reset if required
if not self.isValidEnd(self._last_start, self._last_end):
self._last_end = self._end_frame
start = self._last_start if checked else 1
end = self._last_end if checked else self._total_frame
self.setStartAndEnd(start, end)
self.setStartAndEndEnabled(checked)
[docs] def resetRange(self):
"""
Resets 'Start' & 'End'
"""
self._last_start = self._start_frame
self._last_end = self._end_frame
self.setStartAndEnd(self._start_frame, self._end_frame)
enable = self.frame_count != self._total_frame
self.ui.limit_checkbox.setChecked(enable)
self.setStartAndEndEnabled(enable)
[docs] def setStartAndEnd(self, start, end):
"""
Set 'Start' & 'End' with given values, update entry count labels, reset
button, and style.
"""
self.start = start
self.end = end
self.ui.reset_button.setVisible((self.start != self._start_frame) or
(self.end != self._end_frame))
self.updateRangeStyle()
self.updateEntryCountLabels()
[docs] def setStartAndEndEnabled(self, enable):
"""
Enable/disable 'Start' & 'End'
"""
wids = {'start_label', 'start_lineedit', 'end_label', 'end_lineedit'}
for wid in wids:
attr = getattr(self.ui, wid)
attr.setEnabled(enable)
[docs] def startChanged(self):
"""
Triggered when 'Start' changes
"""
self._last_start = self.start
self.handleRangeChange()
[docs] def endChanged(self):
"""
Triggered when 'End' changes
"""
self._last_end = self.end
self.handleRangeChange()
[docs] def handleRangeChange(self):
"""
When 'Start' or 'End' modified, then show reset button,
update entry count labels and range style
"""
self.ui.reset_button.setVisible(True)
self.updateEntryCountLabels()
self.updateExportButton()
self.updateRangeStyle()
[docs] def updateRangeStyle(self):
"""
Updates 'Start' & 'End' line edit style based on the values
"""
prop_name = "invalid"
self.setWidgetStylePropertyValue(
self.ui.start_lineedit, prop_name,
"false" if self.isValidStart() else "true")
self.setWidgetStylePropertyValue(
self.ui.end_lineedit, prop_name,
"false" if self.isValidEnd() else "true")
[docs] def help(self):
"""
Shows 'Export Structures' dialog help
"""
if maestro:
maestro.command("helptopic TRAJECTORY_EXPORT_STRUCTURES")
maestro.command("showpanel help")