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.
FILE_SELECTOR_CLASS = file_selector.FileSelector
[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 self._file_formats is None:
self._file_formats = [fileutils.MAESTRO]
elif 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 = self.FILE_SELECTOR_CLASS(
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.layoutFileSelector()
[docs] def layoutFileSelector(self):
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]
[docs]class SlimStructureSelector(StructureSelector):
"""
This is a slim version of Structure Selector that spans only one row
of widgets. File selection is handled by a single "Browse..." button,
and there is no widget for showing the file path. It's up to individual
panels to render the loaded file or entries somewhere in its UI.
To the right of the widget, a "Load" button is present, which gets
hidden for File input source. Clicking it causes inputLoaded signal
to be emitted. This signal is also emitted when new file is browsed.
To rename the load button:
SlimStructureSelector.load_btn.textText(<name>)
"""
model_class = StructureSelectorModel
inputLoaded = QtCore.pyqtSignal()
FILE_SELECTOR_CLASS = file_selector.SlimFileSelector
[docs] def initSetUp(self):
super().initSetUp()
self.load_btn = QtWidgets.QPushButton("Load")
self.load_btn.clicked.connect(self.inputLoaded)
if InputSource.File in self._sources:
self.file_selector.fileSelectionChanged.connect(self.inputLoaded)
[docs] def initLayOut(self):
super().initLayOut()
self.source_layout.insertWidget(2, self.load_btn)
[docs] def layoutFileSelector(self):
self.source_layout.insertWidget(3, self.file_selector)
def _updateFileWidgetsVisible(self):
file_source = self.model.current_source is InputSource.File
self.file_selector.setVisible(file_source)
self.load_btn.setVisible(not file_source)
# 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).run(blocking=True)
application.start_application(_main)