Source code for schrodinger.application.livedesign.ld_tasks

"""
Task class used for LiveDesign import GUIs.

Copyright Schrödinger, LLC. All rights reserved.
"""
import typing

from schrodinger import structure
from schrodinger.application.livedesign import export_models
from schrodinger.models import parameters
from schrodinger.tasks import tasks


class ImportTaskInput(export_models.LDClientModelMixin,
                      parameters.CompoundParam):
    ld_client: object
    lr_id: str
    lr_columns: list  # Contains ld_models.Column


class ImportTask(tasks.ThreadFunctionTask):
    """
    Import specified data from a LiveReport
    """

    input: ImportTaskInput

    class Output(parameters.CompoundParam):
        _sts_file_str: str
        imported_structures: typing.List[structure.Structure]

    def mainFunction(self):
        """
        Imports structures from maestro.
        This logic will be run in a worker thread and must not
        access thread-unsafe libraries, including structure.Structure.
        """
        ld_client = self.input.ld_client
        lr_id = self.input.lr_id
        lr_columns = self.input.lr_columns
        col_ids = [col.id for col in lr_columns]
        file_str = ld_client.export_to_maestro(lr_id, col_ids)
        self.output._sts_file_str = file_str

    @tasks.postprocessor
    def postProcessStructures(self):
        """
        Read in structures in a post processor to avoid working
        with structures in a non-thread safe way
        """
        file_str = str(self.output._sts_file_str, encoding='utf-8')
        with structure.StructureReader.fromString(file_str) as reader:
            sts = list(reader)
        self.output.imported_structures = sts