import collections
import enum
import typing
from typing import List
from schrodinger import structure
from schrodinger.models import parameters
from schrodinger.utils.fileutils import force_remove
from . import constants
from . import data_classes
from . import entry_types as ets
from . import ld_utils
from . import login
from . import upload_utils
PROPNAME_COMPOUND_ID = constants.PROPNAME_COMPOUND_ID
PROPNAME_CORP_ID = constants.PROPNAME_CORP_ID
PROPNAME_IMPORT_ENTITY_ID = constants.PROPNAME_IMPORT_ENTITY_ID
LD_DATA_3D = data_classes.LDData(user_name=constants.USERNAME_3D_DATA,
family_name=constants.FAMNAME_3D_DATA,
requires_3d=True)
RefreshResult = enum.Enum('RefreshResult', ('none', 'success', 'failure'))
[docs]class MatchCompoundsBy(enum.Enum):
structure = 'Structure or Imported Corporate ID'
corp_id = 'Corporate ID'
def __str__(self):
return self.value
[docs]class LDDestination(parameters.CompoundParam):
"""
Parameters specifying the destination of the exported data, both LiveDesign
server and live report.
"""
host: str
proj_id: str
proj_name: str
lr_id: str
lr_name: str
[docs]class DataExportSpec(parameters.CompoundParam):
"""
Abstract specification for uploading data to a LiveDesign server.
"""
data_name: str
ld_model: str
ld_endpoint: str
units: str
option: object
[docs] def addDataToExportTask(self, task):
"""
Update the provided task with data from this specification.
Must be overridden in concrete subclasses.
:param task: an export task
:type task: tasks.ThreadFunctionTask
"""
raise NotImplementedError
[docs]class PropertyExportSpec(DataExportSpec):
"""
Specification for structure property data.
"""
[docs] def addDataToExportTask(self, task):
"""
Update the provided task `prop_dicts` attribute with data from this
specification.
:param task: an export task
:type task: tasks.ThreadFunctionTask
"""
prop_dict = make_prop_dict(units=self.units,
name=self.getName(),
endpoint=self.ld_endpoint,
model=self.ld_model)
task.prop_dicts += [prop_dict]
[docs] def getName(self):
"""
:return: the appropriate value for the "name" field for this property's
property dictionary during export to LiveDesign
:rtype: str
"""
if ld_utils.is_sd_dataname(self.data_name):
# SD files store SD data according to a format different
# from the standard structure property data name; they should be
# formatted as if they were a user name, e.g. "All IDs" rather than
# "s_sd_All_IDs"
property_name = structure.PropertyName(dataname=self.data_name)
return property_name.userName()
else:
return self.data_name
[docs]class Base3DExportSpec(DataExportSpec):
"""
Abstract specification for 3D structure data.
"""
[docs] def addDataToExportTask(self, task, rl_map, corp_id_match_prop=None):
"""
Update the provided task `three_d_export_items` attribute with data
from this specification.
:param task: an export task
:type task: tasks.ThreadFunctionTask
:param rl_map: a receptor-ligand map
:type rl_map: data_classes.ReceptorLigandMap
:param corp_id_match_prop: optionally, a property that stores the
corporate ID that should be used to store these 3D items on LD, if
any
:type corp_id_match_prop: str or NoneType
"""
items = self._prepare3DExportItems(
rl_map, corp_id_match_prop=corp_id_match_prop)
task.input.three_d_export_items.extend(items)
def _prepare3DExportItems(self, rl_map, corp_id_match_prop=None):
"""
Generate a list of 3D export items according to this specification.
Must be overridden in concrete subclasses.
:param rl_map: a receptor-ligand map
:type rl_map: data_classes.ReceptorLigandMap
:param corp_id_match_prop: optionally, a property that stores the
corporate ID that should be used to store these 3D items on LD, if
any
:type corp_id_match_prop: str or NoneType
:return: a list of 3D items for export
:rtype: list[ThreeDExportItem]
"""
raise NotImplementedError
def _getExportKey(self, ligand, rl_group, corp_id_match_prop=None):
"""
If possible, determine the key for the specified ligand.
:param ligand: the ligand to be exported
:type ligand: structure.Structure
:param rl_group: the receptor-ligand group to which the ligand belongs
:type rl_group: data_classes.ReceptorLigandGroup
:param corp_id_match_prop: optionally, a structure property data name
that specifies the source of the corporate ID for each compound
:type corp_id_match_prop: str or NoneType
:return: the appropriate key to use for this ligand, if any
:rtype: str or structure.Structure or NoneType
"""
if rl_group.ligand and ligand != rl_group.ligand:
# If the ligand is not the primary ligand for this RL group, use the
# primary ligand as they key so that it gets grouped with the
# primary ligand
return rl_group.ligand
key = None
if corp_id_match_prop:
# If the user specifies a property from which to obtain the
# corporate ID, get the corporate ID from that property
key = rl_group.ligand.property.get(corp_id_match_prop)
# Corporate ID when available should only be a string value.
key = None if key is None else str(key)
if key is None and ld_utils.st_matches_round_trip_hash(rl_group.ligand):
# If no corporate ID has been specified but the ligand structure was
# downloaded from LD, attempt to match the ligand with its
# source compound on LD
key = rl_group.ligand.property.get(PROPNAME_IMPORT_ENTITY_ID)
return key
[docs]class Standard3DExportSpec(Base3DExportSpec):
"""
Specification for standard 3D export.
"""
def _prepare3DExportItems(self, rl_map, corp_id_match_prop=None):
"""
Generate a list of 3D export items according to this specification.
:param rl_map: a receptor-ligand map
:type rl_map: data_classes.ReceptorLigandMap
:param corp_id_match_prop: optionally, a property that stores the
corporate ID that should be used to store these 3D items on LD, if
any
:type corp_id_match_prop: str or NoneType
:return: a list of 3D items for export
:rtype: list[ThreeDExportItem]
"""
items = []
for rl_group in rl_map.rl_groups:
item = ThreeDExportItem()
item.ligand = rl_group.alt_ligand or rl_group.ligand
item.receptor = rl_group.receptor
key = self._getExportKey(item.ligand, rl_group, corp_id_match_prop)
item.setItemKey(key)
item.three_d_specs = [self]
items.append(item)
return items
[docs]class AttachmentItem(parameters.CompoundParam):
"""
Data class for linking structures with a local and
remote file path for an upload
"""
row_structures: typing.List[structure.Structure]
file_path: str
remote_file_name: str
file_type: str = upload_utils.ATTACHMENT
[docs]class AttachmentData(parameters.CompoundParam):
"""
Data class for storing FFC attachment information.
"""
description: str = None
attachment_items: typing.List[AttachmentItem]
[docs]class FFCExportSpec(DataExportSpec):
"""
Abstract specification for FFC attachment export.
"""
FILE_TYPE: str = upload_utils.ATTACHMENT
DESCRIPTION = NotImplemented
attachment_items: typing.List[AttachmentItem]
def _getColumnName(self):
"""
:return: the name of the column to which this data should be exported
:rtype: str
"""
name = self.ld_model
if self.ld_endpoint:
name += f' {self.ld_endpoint}'
return name
[docs] def getAttachmentData(self, panel_model):
"""
Retrieve and store attachment data from the panel model.
Must be overridden in concrete subclasses.
:param panel_model: the model for the LD Export panel
:type panel_model: ld_export2.ExportModel
"""
raise NotImplementedError
[docs] def addDataToExportTask(self, task):
"""
Update the provided task `attachment_data_map` attribute with data from
this specification.
:param task: an export task
:type task: tasks.ThreadFunctionTask
"""
data = AttachmentData()
data.description = self.DESCRIPTION
attachment_items = self.createAttachmentItems()
data.attachment_items = attachment_items
# Cache attachment items for use in delete later
self.attachment_items = attachment_items
col_name = self._getColumnName()
task.attachment_data_map[col_name] = data
[docs] def createAttachmentItems(self):
"""
Create attachment items for a given task.
These items link structures (rows in the LiveReport) with
the attachment to be uploaded for them.
Must be overridden in concrete subclasses.
:return: a list of attachment items to upload data for
:rtype: list(AttachmentItem)
"""
raise NotImplementedError
[docs] def removeLocalFiles(self):
"""
Remove the files created by this spec, if they still exist.
"""
for attachment_item in self.attachment_items:
force_remove(attachment_item.file_path)
[docs]class ThreeDExportItem(parameters.CompoundParam):
"""
Parameters specifying 3D structure data export.
:ivar key: the identification key for this 3D item. It should either be:
1. a structure, if its corporate ID should be the same as that assigned
to a structure that has not yet been exported
2. a string, if its corporate ID is known
3. `None`, if its corporate ID should be automatically assigned by LD
:vartype key: structure.Structure or str or NoneType
"""
key: object
ligand: structure.Structure = None
receptor: structure.Structure = None
three_d_specs: List[Base3DExportSpec]
[docs] def setItemKey(self, key):
"""
Assign the key to be used to identify where this item should be stored
on a LiveDesign server.
:param key: the identification key for this item
:type key: structure.Structure or str or NoneType
"""
self.key = key
if isinstance(key, str):
self.setLigandCorpID(key)
else:
self.setLigandCorpID(None)
[docs] def setLigandCorpID(self, corp_id):
"""
Assign the corporate ID for this structure.
:param corp_id: the corporate ID, if any
:type corp_id: str or NoneType
"""
if corp_id is None and PROPNAME_CORP_ID in self.ligand.property:
del self.ligand.property[PROPNAME_CORP_ID]
elif corp_id is not None:
self.ligand.property[PROPNAME_CORP_ID] = str(corp_id)
[docs] def getLigandCorpID(self):
"""
:return: the corporate ID for this item
:rtype: str or NoneType
"""
return self.ligand.property.get(PROPNAME_CORP_ID)
[docs] def getLigandCompoundID(self):
"""
:return: the compound ID for this item
:rtype: str or NoneType
"""
return self.ligand.property.get(PROPNAME_COMPOUND_ID)
[docs]class SummaryModel(parameters.CompoundParam):
"""
The model for the summary panel shown the user prior to export.
"""
ld_destination: LDDestination
structures_for_2d_export: List[structure.Structure]
input_summary: str
three_d_export_items: List[ThreeDExportItem]
match_compounds_by: MatchCompoundsBy
property_export_specs: List[PropertyExportSpec]
ffc_export_specs: List[FFCExportSpec]
export_specs: List[DataExportSpec]
publish_data: bool
[docs]class LDClientModelMixin:
"""
Mixin for models that contain a `LDClient` instance named `ld_client`.
"""
[docs] def refreshLDClient(self):
"""
Check whether the stored `LDClient` instance is connected to LiveDesign.
If not, create a new instance and replace the old one if the new one is
connected.
:return: an enum describing the status of the connection;
- `none` if no refresh was required
- `success` if the the `LDClient` instance was replaced
- `failure` if even the new `LDClient` instance was disconnected
:rtype: `RefreshResult`
"""
if self.ld_client is None or not ld_utils.is_connected(self.ld_client):
_, ld_client, _ = login.get_ld_client_and_models()
if ld_client and ld_utils.is_connected(ld_client):
self.ld_client = ld_client
return RefreshResult.success
return RefreshResult.failure
return RefreshResult.none
[docs]class PoseNameEditModel(parameters.CompoundParam):
"""
Model for the Pose Name Edit Panel.
:ivar custom_text: the text of the custom text line edit; this value is
stored temporarily while the panel is open, and will be copied to
`custom_text_final` if the user accepts the panel
:ivar include_property: the check state of the "include property" checkbox
:ivar property_name: the structure property (if any) selected as part of
the custom pose name; this value is stored temporarily while the panel
is open, and will be copied to `property_name_final` if the user
accepts the panel
:ivar property_user_name: the text of the structure property label
:ivar example_prop_string: the text of the example property
:ivar example_name: the text of the example pose name label
:ivar entry_data: the system entry data for the panel
:ivar custom_text_final: the custom text accepted by the user
:ivar property_name_final: the structure property (if any) accepted by the
user
"""
custom_text: str
include_property: bool
property_name: structure.PropertyName = None
property_user_name: str = '(not defined)'
example_prop_string: str = None
example_name: str
entry_data: ets.BaseEntryData = None
custom_text_final: str
property_name_final: structure.PropertyName = None
[docs]class AttachmentTaskOutput(parameters.CompoundParam):
"""
Output model for FFC attachment export task.
"""
num_success: int
num_failure: int
[docs]class FileBatch(parameters.CompoundParam):
"""
Data class for storing file paths necessary for standard LD export (v8.6-).
"""
map_file_path: str = None
sdf_file_path: str = None
three_d_file_path: str = None
[docs]def make_prop_dict(units='', name='', endpoint='', model=''):
"""
Return a dictionary formatted to represent a structure property. Formatting
should match the dictionary formatting required by the `properties` argument
of `LDClient.start_export_assay_and_pose_data()`.
:param units: unit system used by this property
:type units: str
:param name: name for this property; either the data name or user name
:type name: str
:param endpoint: the LiveDesign endpoint
:type endpoint: str
:param model: the user-specified portion of the name for the column under
which this property will be stored after export to LiveDesign
:type model: str
"""
prop_dict = collections.OrderedDict()
# SD files get rid of the 's_m_title' field altogether, storing
# the title outside of the property dictionary.
if name == constants.PROPNAME_TITLE:
name = constants.PROPNAME_SD_TITLE
prop_dict[constants.LD_PROP_UNITS] = units
prop_dict[constants.LD_PROP_NAME] = name
prop_dict[constants.LD_PROP_ENDPOINT] = endpoint
prop_dict[constants.LD_PROP_MODEL] = model
return prop_dict