read_sdf()

Complexity

Conditions

18

Size

Total Lines

97

Duplication

Lines	0
Ratio	0 %

Code Coverage

Tests	25
CRAP Score	21.3891

Importance

Changes	8
Bugs	0	Features	3

#! /usr/bin/env python
def read_sdf(sdf, error_bad_mol=False, warn_bad_mol=True, nmols=None,

Too many arguments (8/5)

This function exceeds the maximum number of variables (18/15).

             skipmols=None, skipfooter=None, read_props=True, mol_props=False,
             *args, **kwargs):

    """Read an sdf file into a `pd.DataFrame`.

    The function wraps the RDKit `ForwardSDMolSupplier` object.

    Args:
        sdf (str or file-like):
            The location of data to load as a file path, or a file-like object.
        error_bad_mol (bool):
            Whether an error should be raised if a molecule fails to parse.
            Default is False.
        warn_bad_mol (bool):
            Whether a warning should be output if a molecule fails to parse.
            Default is True.
        nmols (int):
            The number of molecules to read. If `None`, read all molecules.
            Default is `None`.
        skipmols (int):
            The number of molecules to skip at start.
            Default is `0`.
        skipfooter (int):
            The number of molecules to skip from the end.
            Default is `0`.
        read_props (bool):
            Whether to read the properties into the data frame.
            Default is `True`.
        mol_props (bool):
            Whether to keep properties in the molecule dictionary after they
            are extracted to the DataFrame.
            Default is `False`.
        args, kwargs:
            Arguments will be passed to RDKit ForwardSDMolSupplier.

    Returns:
        pandas.DataFrame:
            The loaded data frame, with Mols supplied in the `structure` field.

    See also:
        rdkit.Chem.SDForwardMolSupplier
        skchem.read_smiles
    """

    # nmols is actually the index to cutoff.  If we skip some at start, we need
    # to add this number
    if skipmols:
        nmols += skipmols

    if isinstance(sdf, str):
        sdf = open(sdf, 'rb')  # use read bytes for python 3 compatibility

    # use the suppression context manager to not pollute our stdout with rdkit
    # errors and warnings.
    # perhaps this should be captured better by Mol etc.
    with Suppressor():

        mol_supp = Chem.ForwardSDMolSupplier(sdf, *args, **kwargs)

        mols = []

        # single loop through sdf
        for i, mol in enumerate(mol_supp):

            if skipmols and i < skipmols:
                continue

            if nmols and i >= nmols:
                break

            if mol is None:
                msg = 'Molecule {} could not be decoded.'.format(i + 1)
                if error_bad_mol:
                    raise ValueError(msg)
                elif warn_bad_mol:
                    warnings.warn(msg)
                continue

            mols.append(Mol(mol))

        if skipfooter:
            mols = mols[:-skipfooter]

    idx = pd.Index((m.name for m in mols), name='batch')
    data = pd.DataFrame(mols, columns=['structure'])

    if read_props:
        props = pd.DataFrame([{k: v for (k, v) in mol.props.items()}
                              for mol in mols])
        data = pd.concat([data, props], axis=1)
        # now we have extracted the props, we can delete if required
        if not mol_props:
            data.apply(_drop_props, axis=1)

    data.index = idx
    return squeeze(data, axis=1)