Source code for schrodinger.application.steps.transformers

from rdkit import Chem

from schrodinger.infra import mmerr
from schrodinger.structutils.build import desalt_structure
from schrodinger.structutils.build import neutralize_structure
from schrodinger.utils import license

from . import utils
from .basesteps import MolMapStep

STANDARDIZER_MMERR_LEVEL = mmerr.MMERR_OFF


class Standardizer(MolMapStep):
    """
    Creates a coordinate-less molecule that is neutralized and desalted.

    This involves a conversion to a `schrodinger.structure.Structure`.
    """

    def getLicenseRequirements(self):
        return {license.LIGPREP_MAIN: 1}

    def mapFunction(self, mol):
        st = utils.mol_to_structure(mol, self, generate_coordinates=False)
        if st is not None:
            with mmerr.Level(STANDARDIZER_MMERR_LEVEL):
                st = desalt_structure(neutralize_structure(st))
            mol = utils.structure_to_mol(st, self, mol)
            if mol is not None:
                yield mol


class StereoChemistryRemover(MolMapStep):
    """
    Removes the stereo chemistry from a molecule.
    """

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