Source code for schrodinger.pipeline.stages.example

from schrodinger import structure
from schrodinger.pipeline import pipeio
from schrodinger.pipeline import stage


class ExampleStage(stage.Stage):

    def __init__(self, *args, **kwargs):
        stage.Stage.__init__(self, *args, **kwargs)

        specs = """
        KEEP_STS = integer(default=10)
        """
        stage.Stage.__init__(self, specs=specs, *args, **kwargs)

        self.addExpectedInput(1, "structures", required=True)
        self.addExpectedOutput(1, "structures", always=True)

    def operate(self):
        keep_sts = self['KEEP_STS']

        # Read the input pin:
        input_files = self.getInput(1).getFiles()

        # Prepare the output writer:
        out_file = self.getOutputName(1) + ".maegz"
        writer = structure.StructureWriter(out_file)

        st_num = 0
        for st in structure.MultiFileStructureReader(input_files):
            writer.append(st)
            st_num += 1

            if st_num == keep_sts:
                break

        writer.close()  # close the structure writer

        self.info("Total %i structures kept" % st_num)

        self.setOutput(1, pipeio.Structures([out_file], st_num))


# EOF