import enum
import os
from typing import List
import inflect
import schrodinger
from schrodinger import structure
from schrodinger.maestro_utils import maestro_sync
from schrodinger.models import mappers
from schrodinger.models import parameters
from schrodinger.project import utils as project_utils
from schrodinger.project.utils import get_PT
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtWidgets
from schrodinger.ui.qt import basewidgets
from schrodinger.ui.qt import file_selector
from schrodinger.ui.qt import filedialog
from schrodinger.ui.qt.appframework2 import application
from schrodinger.ui.qt.appframework2 import validation
from schrodinger.ui.qt.mapperwidgets import MappableComboBox
from schrodinger.utils import fileutils
maestro = schrodinger.get_maestro()
ERR_NO_FILE = 'No file specified as a structure source'
ERR_NO_SEL = 'No Project Table entries are selected'
ERR_NO_INC = 'No Project Table entries are included in the Workspace'
ERR_ONE_INC = 'Only one entry must be included in the Workspace'
ERR_FILE_EXIST = 'The specified file does not exist'
# TODO: Add Workspace option that allows multiple entries to be merged
# into a single structure.
INCLUDED_SOURCES = (InputSource.IncludedEntry, InputSource.IncludedEntries)
MAE_DEPENDENT_SOURCES = (InputSource.SelectedEntries, *INCLUDED_SOURCES)
[docs]class StructureSelectorMaestroSync(maestro_sync.BaseMaestroSync,
QtCore.QObject):
"""
:ivar projectUpdated: a signal to indicate that the included or selected
entries have changed. Emitted with whether Workspace inclusion has
changed.
"""
inclusionChanged = QtCore.pyqtSignal()
selectionChanged = QtCore.pyqtSignal()
[docs] def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._previously_included_eids = None
self._previously_selected_eids = None
self.addProjectUpdateCallback(self._onProjectUpdated)
[docs] def getNumIncludedEntries(self):
pt = get_PT()
return len(pt.included_rows)
def _getIncludedEntryIDs(self):
return set(project_utils.get_included_entry_ids())
[docs] def getNumSelectedEntries(self):
pt = get_PT()
return pt.getSelectedRowTotal()
def _getSelectedEntryIDs(self):
return set(project_utils.get_selected_entry_ids())
[docs] def getIncludedStructures(self):
return project_utils.get_included_structures()
[docs] def getSelectedStructures(self):
return project_utils.get_selected_structures()
[docs] def inclusionAndSelectionAreIdentical(self):
return self._getIncludedEntryIDs() == self._getSelectedEntryIDs()
def _onProjectUpdated(self):
included_eids = self._getIncludedEntryIDs()
inclusion_changed = (included_eids != self._previously_included_eids)
self._previously_included_eids = included_eids
if inclusion_changed:
self.inclusionChanged.emit()
selected_eids = self._getSelectedEntryIDs()
selection_changed = (selected_eids != self._previously_selected_eids)
self._previously_selected_eids = selected_eids
if selection_changed:
self.selectionChanged.emit()
[docs]class StructureSelectorModel(parameters.CompoundParam):
"""
Note presence of both `file_paths` and `file_path`. These are both
included because it's unclear whether the file selector will be initialized
with our without support for multiple files.
If multiple files are supported, `StructureSelector.getFilePaths()` should
be used to access them.
If multiple files are not supported, either
`StructureSelector.getFilePath()` or `getFilePaths()` may be used to acess
the selected file.
"""
# NOTE: Outside of Maestro, only File source is shown
sources: List[InputSource] = [InputSource.SelectedEntries, InputSource.File]
current_source: InputSource = None # by default first one in list
source_lbl_text: str = 'Use structures from:'
file_paths: List[str]
file_path: str
[docs]class StructureSelector(mappers.MapperMixin, basewidgets.BaseWidget):
"""
Widget to extract structures from maestro and/or files.
"""
model_class = StructureSelectorModel
# Emitted when input source is changed (e.g. workspace vs file):
sourceChanged = QtCore.pyqtSignal(InputSource)
# Emitted when input structures are changed:
inputChanged = QtCore.pyqtSignal()
# TODO implement inputChangedExtraTracking signal.
[docs] def __init__(self,
parent,
sources=None,
default_source=None,
file_formats=None,
support_multiple_files=False,
initial_dir=None):
"""
Initialize the StructureSelector.
:param parent: Parent widget.
:type parent: QWidget
:param sources: Supported input sources. Default is SelectedEntries
and File.
:type sources: List(InputSource)
:param default_source: Default source. Default is the first source
in the "sources" list.
:param file_formats: Supported file formats; see fileutls.py module.
E.g. file_formats=[fileutils.MAESTRO].
:type file_formats: list
:param support_multiple_files: Whether or not to allow the user to
select multiple files at once from the file dialog.
:type support_multiple_files: bool
:param initial_dir: Initial directory. Default is CWD.
:type initial_dir: str
"""
if sources is None:
sources = [InputSource.SelectedEntries, InputSource.File]
if not maestro:
sources = [s for s in sources if s not in MAE_DEPENDENT_SOURCES]
if not sources:
msg = 'Must specify at least one input source.'
if not maestro:
msg += (' Maestro-dependent input sources '
f'({", ".join(s.name for s in MAE_DEPENDENT_SOURCES)}) '
'are ignored outside of Maestro.')
raise ValueError(msg)
if (InputSource.IncludedEntries in sources and
InputSource.IncludedEntry in sources):
raise ValueError("Sources IncludedEntries and IncludedEntry are "
"mutually exclusive.")
self._sources = sources
self._default_source = default_source
self._file_formats = file_formats
self._support_multiple_files = support_multiple_files
self._initial_dir = initial_dir
super().__init__(parent)
[docs] def initSetUp(self):
super().initSetUp()
self.source_lbl = QtWidgets.QLabel() # text will be taken from model
self.source_combo = MappableComboBox(self)
for source in self._sources:
text = source.value
self.source_combo.addItem(text, source)
# TODO: Hide combo if only one option is available
if InputSource.File in self._sources:
if not isinstance(self._file_formats, list):
msg = ('Expected a list of supported structure file formats'
f' but instead got {type(self._file_formats)}.')
raise TypeError(msg)
filter_str = filedialog.filter_string_from_formats(
self._file_formats)
self.file_selector = file_selector.FileSelector(
parent=self,
filter_str=filter_str,
support_multiple_files=self._support_multiple_files,
initial_dir=self._initial_dir)
self.file_selector.fileSelectionChanged.connect(self.inputChanged)
else:
if self._file_formats:
msg = ('InputSource.File is not specified as an input source,'
' but allowed file formats have been specified. Either'
' remove the `file_formats` argument from the'
' constructor or include InputSource.File in the'
' `sources` argument.')
raise ValueError(msg)
self._previous_input_source = None
self._previous_valid = True
if maestro:
self._maestro_sync = StructureSelectorMaestroSync()
self._maestro_sync.selectionChanged.connect(
self._onSelectionChanged)
self._maestro_sync.inclusionChanged.connect(
self._onInclusionChanged)
[docs] def initLayOut(self):
super().initLayOut()
self.createSourceLayout()
self.widget_layout.addLayout(self.source_layout)
if InputSource.File in self._sources:
self.widget_layout.addWidget(self.file_selector)
[docs] def createSourceLayout(self):
self.source_layout = QtWidgets.QHBoxLayout()
self.source_layout.addWidget(self.source_lbl)
self.source_layout.addWidget(self.source_combo)
self.source_layout.addStretch()
[docs] def showEvent(self, event):
"""
Enable Maestro callbacks when the panel containing the StructureSelector
is shown.
:type event: QShowEvent
:param event: The QEvent object generated by this event.
:return: The return value of the base class showEvent() method.
"""
value = super().showEvent(event)
if maestro:
self._maestro_sync.setCallbacksActive(True)
return value
[docs] def hideEvent(self, event):
"""
Disable Maestro callbacks when the panel containing the StructureSelector
is hidden.
:type event: QHideEvent
:param event: The QEvent object generated by this event.
:return: The return value of the base class hideEvent() method.
"""
value = super().hideEvent(event)
if maestro:
self._maestro_sync.setCallbacksActive(False)
return value
[docs] def defineMappings(self):
M = self.model_class
mappings = [
(self.source_combo, M.current_source),
(self.source_lbl, M.source_lbl_text),
]
if InputSource.File in self._sources:
mappings.append((self._updateFileWidgetsVisible, M.current_source))
if self._support_multiple_files:
mappings.append((self.file_selector, M.file_paths))
else:
mappings.append((self.file_selector, M.file_path))
return mappings
[docs] def getSignalsAndSlots(self, model):
return [
(model.current_sourceChanged, self.sourceChanged),
(model.current_sourceChanged, self._onCurrentSourceChanged),
]
[docs] def getFilePath(self):
"""
Getter for `self.model.file_path`. Should be used only when certain
that multiple file selection is not supported.
:return: The currently selected input file.
:rtype: str
"""
if self._support_multiple_files:
msg = 'Use getFilePaths() when multiple input files are supported'
raise RuntimeError(msg)
return self.model.file_path
[docs] def getFilePaths(self):
"""
Getter for `self.model.file_paths` or `self.model.file_path`. Can
be used when supporting single or multiple files.
:return: The currently selected input file(s).
:rtype: list(str)
"""
if self._support_multiple_files:
return self.model.file_paths
return [self.model.file_path] if self.model.file_path else []
[docs] def setSourceLabelText(self, text):
self.model.source_lbl_text = text
[docs] def setFileLabelText(self, text):
self.file_selector.setFileLabelText(text)
@validation.validator()
def widgetStateIsValid(self):
"""
Validate that the StructureSelector is in a consistent and complete
state.
:return: True if the widget is in a valid state. False otherwise.
:rtype: bool
"""
source = self.model.current_source
if source == InputSource.File:
files = self.getFilePaths()
if not files:
return False, ERR_NO_FILE
for fname in files:
if not os.path.isfile(fname):
return False, f'{ERR_FILE_EXIST}: {fname}'
if source == InputSource.SelectedEntries:
if self._maestro_sync.getNumSelectedEntries() == 0:
return False, ERR_NO_SEL
if source in INCLUDED_SOURCES:
num_inc = self._maestro_sync.getNumIncludedEntries()
if num_inc == 0:
return False, ERR_NO_INC
if source == InputSource.IncludedEntry and num_inc > 1:
return False, ERR_ONE_INC
return True
[docs] def structures(self):
result = self.widgetStateIsValid()
if not result:
raise ValueError(result.message)
inputsource = self.model.current_source
if inputsource is InputSource.File:
for fname in self.getFilePaths():
with structure.StructureReader(fname) as rdr:
yield from rdr
elif inputsource is InputSource.SelectedEntries:
yield from self._maestro_sync.getSelectedStructures()
elif inputsource in INCLUDED_SOURCES:
yield from self._maestro_sync.getIncludedStructures()
else:
raise RuntimeError
[docs] def countStructures(self):
"""
Return the number of structures specified in the selector.
:return: The number of structures; or 0 on invalid state.
:rtype: int
"""
if not self.widgetStateIsValid():
return 0
inputsource = self.model.current_source
if inputsource is InputSource.File:
return sum(
structure.count_structures(fname)
for fname in self.getFilePaths())
elif inputsource is InputSource.SelectedEntries:
return self._maestro_sync.getNumSelectedEntries()
elif inputsource in INCLUDED_SOURCES:
return self._maestro_sync.getNumIncludedEntries()
else:
raise RuntimeError
[docs] def writeFile(self, filename):
"""
Writes the selected structures to the given file path.
:param filename: File path
"""
if not self.widgetStateIsValid():
raise ValueError
structure.write_cts(self.structures(), filename)
def _storeValid(self):
self._previous_valid = bool(self.widgetStateIsValid())
def _onCurrentSourceChanged(self, input_source):
"""
Called when the input source is modified.
"""
prev_valid = self._previous_valid
self._storeValid()
prev_input_source = self._previous_input_source
self._previous_input_source = input_source
if input_source in MAE_DEPENDENT_SOURCES and prev_input_source in MAE_DEPENDENT_SOURCES:
# If one source is included and one source is selected, skip signal
# if same entries are included and selected and validity is the same
# (The validity check is needed because
# IncludedEntry is invalid if multiple entries are included)
if (prev_valid == self._previous_valid
) and self._maestro_sync.inclusionAndSelectionAreIdentical():
return
self.inputChanged.emit()
def _onInclusionChanged(self):
"""
Called when Workspace inclusion is changed.
"""
combo = self.source_combo
Included = InputSource.IncludedEntries
if Included in self._sources:
num_included_entries = self._maestro_sync.getNumIncludedEntries()
entry_text = inflect.engine().no('entry', num_included_entries)
text = f"Workspace ({entry_text})"
included_idx = combo.findData(Included)
combo.setItemText(included_idx, text)
if self.model.current_source in INCLUDED_SOURCES:
self._storeValid()
self.inputChanged.emit()
def _onSelectionChanged(self):
"""
Called when Project Table selection is changed.
"""
combo = self.source_combo
Selected = InputSource.SelectedEntries
if Selected in self._sources:
num_selected_entries = self._maestro_sync.getNumSelectedEntries()
text = "Project Table (%i selected)" % num_selected_entries
selected_idx = combo.findData(Selected)
combo.setItemText(selected_idx, text)
if self.model.current_source == Selected:
self._storeValid()
self.inputChanged.emit()
def _updateFileWidgetsVisible(self):
"""
Called when input source is changed - shows or hides the file related
options as needed.
"""
files_visible = self.model.current_source is InputSource.File
fs = self.file_selector
file_widgets = (fs.file_le, fs.file_lbl, fs.file_browse_btn)
for wdg in file_widgets:
wdg.setVisible(files_visible)
[docs] def initSetDefaults(self):
"""
Select the default source.
"""
super().initSetDefaults()
if maestro and self._default_source:
self.model.current_source = self._default_source
else:
self.model.current_source = self._sources[0]
# For testing purposes only: #
[docs]def panel():
"""
For testing StructureSelector widget within Maestro.
"""
inp_sel = StructureSelector(None)
inp_sel.show()
if __name__ == "__main__":
# For testing StructureSelector widget outside of Maestro.
def _main():
StructureSelector(parent=None,
file_formats=[fileutils.MAESTRO]).run(blocking=True)
application.start_application(_main)