import uuid
from contextlib import contextmanager
from schrodinger import get_maestro
from schrodinger import project
from schrodinger.infra import mm
from schrodinger.project import ProjectException
from schrodinger.ui.qt import utils as qt_utils
maestro = get_maestro()
PROPNAME_MARK = 'b_m_Mark'
NOT_IN_WORKSPACE = project.NOT_IN_WORKSPACE
IN_WORKSPACE = project.IN_WORKSPACE
LOCKED_IN_WORKSPACE = project.LOCKED_IN_WORKSPACE
PROP_TRJ = mm.M2IO_DATA_TRAJECTORY_FILE
[docs]def get_PT():
"""
Safe getter for the project instance that swallows exceptions that occur if
the project has been closed.
Using this function avoids unnecessary tracebacks that occur if the project
is requested in the brief period after the project has been closed before a
new one is created.
:return: a project instance, if a project is available in Maestro
:rtype: project.Project or NoneType
"""
if not maestro:
return None
try:
return maestro.project_table_get()
except ProjectException: # project was closed
return None
[docs]def set_entries_pinned(entry_ids, pin):
"""
Pin or unpin entries from the workspace.
:param entry_ids: a list of entry IDs
:type entry_ids: list(str)
:param pin: whether to pin or unpin the specified entries to the workspace
:type pin: bool
"""
pt = maestro.project_table_get()
for entry_id in entry_ids:
row = pt.getRow(entry_id)
if pin:
row.in_workspace = LOCKED_IN_WORKSPACE
elif row.in_workspace == LOCKED_IN_WORKSPACE:
row.in_workspace = IN_WORKSPACE
[docs]def set_entry_pinned(entry_id, pin):
"""
Pin or unpin an entry from the workspace.
:param entry_id: an entry ID
:type entry_id: str
:param pin: whether to pin or unpin the specified entry to the workspace
:type pin: bool
"""
set_entries_pinned([entry_id], pin)
[docs]def set_entries_included(entry_ids, include, override_pin=False):
"""
For the specified entries, either include them in or exclude them from the
workspace.
:param entry_ids: a list of entry IDs
:type entry_ids: list(str)
:param include: whether to include (`True`) or exclude (`False`)
:type include: bool
:param override_pin: if `False`, ignore pinned entries. If `True`, exclude
entries that are set to be excluded even if they are pinned. If pinned
entries are set to be included, leave them as pinned (because they are
also included as pinned entries)
:type override_pin: bool
"""
pt = maestro.project_table_get()
for entry_id in entry_ids:
row = pt.getRow(entry_id)
is_pinned = row.in_workspace == LOCKED_IN_WORKSPACE
if not override_pin and is_pinned:
continue
elif include and not is_pinned:
# If an entry is already pinned, setting it to be IN_WORKSPACE will
# un-pin it
row.in_workspace = IN_WORKSPACE
elif not include:
row.in_workspace = NOT_IN_WORKSPACE
[docs]def set_entry_included(entry_id, include, override_pin=False):
"""
For the specified entry, either include it in or exclude it from the
workspace.
:param entry_id: an entry ID
:type entry_id: str
:param include: whether to include (`True`) or exclude (`False`)
:type include: bool
:param override_pin: if `False`, ignore pinned entries. If `True`, exclude
entries that are set to be excluded even if they are pinned. If pinned
entries are set to be included, leave them as pinned (because they are
also included as pinned entries)
:type override_pin: bool
"""
set_entries_included([entry_id], include, override_pin)
[docs]def entry_is_pinned(entry_id):
"""
Return whether the specified entry is pinned in the workspace.
:param entry_id: the entry ID of a row in the project
:type entry_id: str
:return: whether the entry is pinned in the workspace
:rtype: bool
"""
row = get_row(entry_id)
return row.in_workspace == LOCKED_IN_WORKSPACE
[docs]def entry_is_included(entry_id):
"""
Return whether the specified entry is included in the workspace.
:param entry_id: the entry ID of a row in the project
:type entry_id: str
:return: whether the entry is included in the workspace
:rtype: bool
"""
row = get_row(entry_id)
return row.in_workspace in (IN_WORKSPACE, LOCKED_IN_WORKSPACE)
[docs]def entry_is_selected(entry_id):
"""
Return whether the specified entry is selected in the project.
:param entry_id: the entry ID of a row in the project
:type entry_id: str
:return: whether the entry is selected in the project
:rtype: bool
"""
row = get_row(entry_id)
return row.is_selected
[docs]def set_entry_locked(entry_id, lock):
"""
Set the supplied project table entry to be "locked" or "unlocked."
These operations are grouped because they are both performed when locking
an entry in the visual interface.
:param entry_id: the entry ID of a row in the project
:type entry_id: str
:param lock: whether to lock or unlock the specified entry
:type lock: bool
"""
row = get_row(entry_id)
row.read_only = lock
row.deletable = not lock
[docs]def entry_is_locked(entry_id):
"""
Return whether the supplied project entry is read-only and not deletable.
These properties are both tested because both values are set when locking an
entry in the visual interface.
:param entry_id: the entry ID of a row in the project
:type entry_id: str
:return: whether the entry is locked
:rtype: bool
"""
row = get_row(entry_id)
return row.read_only and not row.deletable
[docs]def remove_entries(entry_ids):
"""
Cleanly remove the specified entries from the project. If an entry cannot be
found, do not raise an exception.
:param entry_id: a list of entry IDs
:type entry_id: list(str)
"""
if not entry_ids:
return
pt = maestro.project_table_get()
# Unlock rows and make them deletable
valid_entry_ids = set()
for entry_id in entry_ids:
row = pt.getRow(entry_id)
if row is None:
continue
set_entry_locked(entry_id, False)
valid_entry_ids.add(entry_id)
with qt_utils.suppress_signals(pt.project_model):
if valid_entry_ids:
entries_str = ', '.join(valid_entry_ids)
maestro.command('entrywsexclude entry_id "{0}"'.format(entries_str))
for entry_id in valid_entry_ids:
pt.deleteRow(entry_id)
pt.update()
[docs]def create_subgroup(entry_ids, group_name, subgroup_title):
"""
Create a subgroup of the specified project group, and place the supplied
entries into it.
:param entry_ids: entry IDs for project entries to be placed into the new
subgroup
:type entry_ids: list(str)
:param group_name: name of the group to which this new subgroup will belong
:type group_name: str
:param subgroup_title: the title of the new subgroup
:type subgroup_title: str
:return: the new group object
:rtype: project.EntryGroup
"""
pt = maestro.project_table_get()
with preserve_selection(pt):
pt.selectRows(entry_ids=entry_ids)
cmd = ('entrygroupcreatewithselected "{0}" autodetectparentgroup=false'
' parentgroup="{1}"').format(subgroup_title, group_name)
maestro.command(cmd)
entry_id = next(iter(entry_ids))
return pt.getRow(entry_id).group
[docs]def move_group(group, row_number):
"""
Move all of the entries in the specified group to be below a specified row
number in the project, while still retaining their group identity.
:param group: a non-empty entry group
:type group: project.EntryGroup
:param row_number: the number of a row in the project: one that designates
order within the project, not an entry ID
:type row_number: int
"""
pt = maestro.project_table_get()
entry_ids = {row.entry_id for row in group.all_rows}
group_name = group.name
parent_group = group.getParentGroup()
parent_group_name = None
if parent_group is not None:
parent_group_name = parent_group.name
with preserve_selection(pt):
pt.selectRows(entry_ids=entry_ids)
# Move all group entries to the desired position
cmd = f'entrymoveselection {row_number} above_row=false'
maestro.command(cmd)
# Re-create the group with the same entries (the above command will have
# moved all entries out of the group)
eid_str = ', '.join(entry_ids)
cmd = f'entrygroupcreate "{group_name}" entry_id "{eid_str}"'
if parent_group_name is not None:
cmd += f' parentgroup="{parent_group_name}"'
maestro.command(cmd)
[docs]def get_rows(entry_ids):
"""
:param entry_id: a list of entry IDs
:type entry_id: list(str)
:return: a list of either project rows corresponding to the provided entry
IDs, if they exist, or `None` otherwise
:rtype: list(project.ProjectRow or None)
"""
# TODO modify this function to raise exception if no such entry exists.
# TODO: Consider returning a generator instead of a list.
pt = maestro.project_table_get()
return [pt.getRow(entry_id) for entry_id in entry_ids]
[docs]def get_row(entry_id):
"""
:param entry_id: an entry ID
:type entry_id: str
:return: row corresponding to the entry ID, if it exists
:rtype: project.ProjectRow or None
"""
# TODO modify this function to raise exception if no such entry exists.
rows = get_rows([entry_id])
return rows[0]
[docs]def get_structures_for_entry_ids(entry_ids, copy=True, props=True, pt=None):
"""
Iterate over structures for the given entry IDs.
"""
if not entry_ids:
return
if maestro:
# if not running in unit tests
maestro.project_table_synchronize()
if pt is None:
pt = maestro.project_table_get()
for eid in entry_ids:
yield pt[eid].getStructure(workspace_sync=False, copy=copy, props=props)
[docs]def get_included_structures(copy=True, props=True, pt=None):
"""
Iterate over structures for entries that are included n the Workspace.
"""
if maestro:
# if not running in unit tests
maestro.project_table_synchronize()
if pt is None:
pt = maestro.project_table_get()
for row in pt.included_rows:
yield row.getStructure(workspace_sync=False, copy=copy, props=props)
[docs]def get_included_entry_ids(pt=None):
"""
Iterate over entry IDs for entries that are included in the Workspace.
:rtype: Iterator[str]
:return: Each iteration yields the entry id of an included project entry
"""
if maestro:
# if not running in unit tests
maestro.project_table_synchronize()
if pt is None:
pt = maestro.project_table_get()
for row in pt.included_rows:
yield row.entry_id
[docs]def get_selected_structures(copy=True, props=True, pt=None):
"""
Iterate over structures for entries that are selected in the Project Table.
"""
if maestro:
# if not running in unit tests
maestro.project_table_synchronize()
if pt is None:
pt = maestro.project_table_get()
for row in pt.selected_rows:
yield row.getStructure(workspace_sync=False, copy=copy, props=True)
[docs]def get_selected_entry_ids(pt=None):
"""
Iterate over entry IDs for entries that are selected in the Project Table.
:rtype: Iterator[str]
:return: Each iteration yields the entry id of an included project entry
"""
if maestro:
# if not running in unit tests
maestro.project_table_synchronize()
if pt is None:
pt = maestro.project_table_get()
for row in pt.selected_rows:
yield row.entry_id
[docs]def get_structure(entry_id):
"""
:param entry_ids: an entry ID
:type entry_ids: str
:return: a structure from the specified entry, if available
:rtype: structure.Structure or None
"""
# TODO modify this function to raise exception if no such entry exists.
pt = maestro.project_table_get()
row = pt.getRow(entry_id)
if row is None:
return None
return row.getStructure()
[docs]@contextmanager
def preserve_selection(pt=None):
"""
Save the selection state of the project, then restore it on exit.
:param pt: optionally, a project instance
:type pt: project.Project or NoneType
"""
pt = pt or maestro.project_table_get()
entry_ids = [row.entry_id for row in pt.selected_rows]
yield
pt.selectRows(entry_ids=entry_ids)
[docs]@contextmanager
def unlock_entries(entry_ids, update_pt=True):
"""
Temporarily unlock the specified entries, if they are locked. Otherwise,
do not modify their lock state. Suppress signals from the project model
during this process.
:param entry_ids: a list of entry IDs
:type entry_ids: list(str)
:param update_pt: whether to update the project table on exiting the
context environment
:type update_pt: bool
"""
pt = maestro.project_table_get()
lock_map = {entry_id: entry_is_locked(entry_id) for entry_id in entry_ids}
with qt_utils.suppress_signals(pt.project_model):
for entry_id in entry_ids:
set_entry_locked(entry_id, False)
yield
with qt_utils.suppress_signals(pt.project_model):
for entry_id, lock in lock_map.items():
if pt.getRow(entry_id) is not None:
set_entry_locked(entry_id, lock)
if update_pt:
pt.update()
[docs]def create_child_group(entry_ids,
parent_group_name,
group_name=None,
before_entry_id=None):
"""
Create a new group using the supplied entry IDs as a child group of the
specified parent group.
:raise ValueError: if no entry IDs are provided
:param entry_ids: a list of entry IDs to put into the new group
:type entry_ids: list(str)
:param parent_group_name: the name of the group that should be the parent of
the group created by this function
:type parent_group_name: str
:param group_name: the name of the group to be created; if not supplied, a
unique name will be randomly generated
:type group_name: str or None
:param before_entry_id: optionally, the entry ID of a project entry above
which the new subgroup should be created
:type before_entry_id: str or NoneType
:return: the new group
:rtype: project.EntryGroup
"""
if not entry_ids:
raise ValueError('Cannot create a group with no entries.')
group_name = group_name or generate_unique_group_name()
esl = 'entry_id ' + ', '.join(entry_ids)
cmd = f'entrygroupcreate {group_name} {esl} parentgroup={parent_group_name}'
if before_entry_id is not None:
cmd += f' positionbeforeentryid={before_entry_id}'
maestro.command(cmd)
return get_entry_group(group_name)
[docs]def get_entry_group(group_name):
"""
Return the group in the project with the specified name.
:param group_name: the name of the desired group
:type group_name: str
:return: the group with the specified name, if available
:rtype: str or None
"""
pt = maestro.project_table_get()
for group in pt.groups:
if group.name == group_name:
return group
[docs]def generate_unique_group_name():
"""
:return: a unique entry group name
:rtype: str
"""
pt = maestro.project_table_get()
extant_group_names = {group.name for group in pt.groups}
group_name = None
while group_name is None or group_name in extant_group_names:
group_name = str(uuid.uuid4())
return group_name
[docs]def get_top_entry(entry_ids):
"""
Given a list of project entry IDs, return the one that appears in the
highest visual row in the project.
:param entry_ids: a list of project entry IDs
:type entry_ids: list(str)
:return: the entry ID corresponding to the "highest" entry in `entry_ids`
:rtype: str
"""
rows = [get_row(entry_id) for entry_id in entry_ids]
rows = sorted(rows, key=lambda row: row.row_number)
return rows[0].entry_id
[docs]def entry_is_marked(entry_id):
"""
Return whether the specified entry is "marked" with the Maestro property
`b_m_Mark`.
:param entry_id: a project entry ID
:type entry_id: str
:return: whether the specified entry has the Mark property, and that the
value of that property is `True`
:rtype: bool
"""
row = get_row(entry_id)
if PROPNAME_MARK in row.property:
return row.property[PROPNAME_MARK]
return False
[docs]def set_entries_marked(entry_ids):
"""
Programmatically mimic the behavior of the Maestro command
"entrymarkincluded", but for an arbitrary group of entries. If all entries
are marked, unmark all entries. Otherwise, mark all entries.
:param entry_ids: a list of entry IDs
:type entry_ids: list(str)
"""
set_marked = not all(entry_is_marked(entry_id) for entry_id in entry_ids)
for entry_id in entry_ids:
row = get_row(entry_id)
row.property[PROPNAME_MARK] = set_marked
[docs]def get_base_entry_group(entry_id):
"""
For the specified entry, return the highest-level parent group to which it
belongs, if possible.
:param entry_id: a project entry ID
:type entry_id: str
:return: the top-level entry group that contains the entry, if the entry
belongs to any group
:rtype: project.EntryGroup or NoneType
"""
row = get_row(entry_id)
group = row.group
if group is None:
return None
parent_group = group.getParentGroup()
while parent_group is not None:
new_parent_group = parent_group.getParentGroup()
group = parent_group
parent_group = new_parent_group
return group
[docs]def move_row(entry_id, row_number):
"""
Move the specified entry row to the specified position in the project.
:param entry_id: an entry ID
:type entry_id: str
:param row_number: the position to which the specified entry should be moved
:type row_number: int
"""
row = get_row(entry_id)
if row.row_number == row_number:
return
with preserve_selection():
row.selectOnly()
cmd = f'entrymoveselection {row_number} above_row=true'
maestro.command(cmd)
[docs]@contextmanager
def modify_row_structure(entry_id=None, row=None):
"""
Context which returns a structure from the row corresponding to the given
entry id and sets the structure back to the row afterwards.
Also preserves the structure's trajectory property.
:param entry_id: entry id of rows to modify
:type entry_id: str
:param `project.ProjectRow` row: The project row for this entry. Either
row or entry_id must be supplied
:return: a structure in a context manager that will set the
structure back onto the project row when control is returned to this
function
:rtype: structure.Structure or NoneType
"""
if entry_id:
row = get_row(entry_id)
if row:
st = row.getStructure()
# Workaround for MAE-41472
row_traj = row.property.get(PROP_TRJ)
if row_traj:
st.property[PROP_TRJ] = row_traj
yield st
row.setStructure(st)
else:
yield None
[docs]def has_valid_wscore_block(row):
"""
Given a ProjectRow instance, return True if the entry has associated
WScore data and a valid receptor; False otherwise.
"""
struc_handle = project.mmproj.mmproj_index_entry_get_ct(
row._project_handle, row._entry_index)
try:
urh = mm.mmct_ct_m2io_get_unrequested_handle(struc_handle)
except mm.MmException:
return False
try:
mm.m2io_goto_block(urh, "m_wsviz_data", 1)
except mm.MmException:
return False
else:
mm.m2io_leave_block(urh)
return True
# There is some feature overlap between the ProjectStructure class and
# modify_row_structure. However, ProjectStructure is safe to use in included
# entry callbacks and modify_row_structure is not (PANEL-16125)
[docs]class ProjectStructure(object):
"""
A Context manager that gets the structure for an entry from the project
table and then optionally puts it back in the PT upon exit
"""
[docs] def __init__(self, row=None, eid=None, modify=True):
"""
Create a ProjectStructure instance
:param `project.ProjectRow` row: The project row for this entry. Either
row or eid must be supplied
:param str eid: The entry ID for this entry. Either row or eid must be
suppled. If the eid is not found in the project, the self.row
attribute will be None.
:param bool modify: If True (default), set the structure back in
the project on exit. If False, do not modify the structure in the
project.
:raise AttributeError: If neither row nor eid are supplied
"""
if eid:
try:
self.row = maestro.project_table_get().getRow(eid)
except project.ProjectException:
# Invalid projects can happen during project close
self.row = None
else:
self.row = row
self.modify = modify
def __enter__(self):
"""
:rtype: `structure.Structure` or None
:return: The structure for this project row or None if the originally
entry ID is no longer valid
"""
if self.row is None:
return None
self.struct = self.row.getStructure()
# This is a workaround for MAE-41472
row_traj = self.row.property.get(PROP_TRJ)
if row_traj:
self.struct.property[PROP_TRJ] = row_traj
return self.struct
def __exit__(self, *args):
if self.row and self.modify:
self.row.setStructure(self.struct)
[docs]def get_trajectory_path(proj, eid):
"""
Method will return None if passed a falsey entry ID.
Return trajectory file path if any, or None.
:param: proj: Project on which to operate.
:type: proj: Project
:param: eid: Entry id associated with the given project.
:type: eid: int or str
"""
return proj.project_model.getTrajectoryPath(int(eid)) if eid else None
[docs]def has_trajectory(proj, eid):
"""
Whether the entry with given entry id has a trajectory associated with it
:param: proj: Project on which to operate.
:type: proj: Project
:param: eid: Entry id associated with the given project.
:type: eid: int or str
"""
return proj.project_model.hasTrajectoryData(int(eid))
[docs]def has_desmond_trajectory(proj, eid):
"""
Whether the entry with given entry id has desmond trajectory
:param: proj: Project on which to operate.
:type: proj: Project
:param: eid: Entry id associated with the given project.
:type: eid: int or str
"""
return proj.project_model.hasDesmondTrajectory(int(eid))
[docs]def has_materials_trajectory(proj, eid):
"""
Whether the entry with given entry id has materials trajectory
:param: proj: Project on which to operate.
:type: proj: Project
:param: eid: Entry id associated with the given project.
:type: eid: int or str
"""
return proj.project_model.hasMaterialsTrajectory(int(eid))
[docs]@contextmanager
def entry_excluded(entry_id: str):
"""
Exclude the given entry temporarily.
:param entry_id: Entry id to be excluded temporarily.
"""
is_included = entry_is_included(entry_id)
is_pinned = entry_is_pinned(entry_id)
if is_included:
set_entry_included(entry_id, False, override_pin=True)
yield
if is_included:
set_entry_included(entry_id, True, override_pin=True)
set_entries_pinned([entry_id], is_pinned)