Source code for schrodinger.application.steps.converters

from rdkit import Chem

from schrodinger import structure
from schrodinger.models import parameters
from schrodinger import stepper
from schrodinger.utils import fileutils

from . import utils
from .dataclasses import MaeInMixin
from .dataclasses import MaeOutMixin
from .dataclasses import MolInMixin
from .dataclasses import MolOutMixin


class SmilesToMolConverter(MolOutMixin, stepper.MapStep):
    """
    Creates a molecule from the first word in the input string, interpreted as
    SMILES.
    """
    Input = str

    def mapFunction(self, string):
        smiles = string.split()[0]
        mol = Chem.MolFromSmiles(smiles)
        # we need to handle errors in the conversion PYAPP-7691
        if mol is not None:
            yield mol


class MolToSmilesConverter(MolInMixin, stepper.MapStep):
    """
    Creates the canonical smiles string for the molecule.
    """
    Output = str

    def mapFunction(self, mol):
        yield Chem.MolToSmiles(mol)


class MolToStructureConverter(MolInMixin, MaeOutMixin, stepper.MapStep):
    """
    Step for converting Chem.Mol to structure.Structure objects.
    """

    class Settings(parameters.CompoundParam):
        generate_coordinates: bool = False

    def mapFunction(self, mol):
        yield utils.mol_to_structure(mol, self,
                                     self.settings.generate_coordinates)


class StructureToMolConverter(MaeInMixin, MolOutMixin, stepper.MapStep):
    """
    Step for converting structure.Structure to Chem.Mol objects.
    """

    def mapFunction(self, st):
        yield utils.structure_to_mol(st, self)


class MaeFiler(MaeInMixin, stepper.ReduceStep):
    """
    A class to write the input `Structure` objects to a Maestro file.
    """
    Output = stepper.StepperFile

    class Settings(parameters.CompoundParam):
        filename: str

    def validateSettings(self):
        filename = self.settings.filename
        if not filename:
            return [stepper.SettingsError(self, 'filename is not defined')]
        if not fileutils.is_maestro_file(filename):
            return [
                stepper.SettingsError(
                    self, f'"{filename}" is not a proper Maestro file name')
            ]
        return []

    def reduceFunction(self, strucs):
        with structure.StructureWriter(self.settings.filename) as writer:
            for struc in strucs:
                writer.append(struc)
        yield stepper.StepperFile(self.settings.filename)