Source code for schrodinger.application.vss.csvsmiles

import csv

from voluptuous import Required
from voluptuous import Schema

from schrodinger.job import jobcontrol
from schrodinger.utils.fileutils import open_maybe_compressed


class CsvSmilesFile:

    SCHEMA = Schema({
        Required('filename'): str,
        Required('smiles_col'): str,
        Required('cid_col'): str,
        'truth_col': str,
    })

    def __init__(self, *, filename, smiles_col, cid_col, truth_col=None):
        self.filename = jobcontrol.get_runtime_path(filename)
        self.cid_col = cid_col
        self.truth_col = truth_col
        self.smiles_col = smiles_col

    def to_dict(self):
        return {k: v for k, v in vars(self).items() if v}

    def validate(self):
        '''
        :return: Validation success and error message.
        :rtype: (bool, str)
        '''
        with open_maybe_compressed(self.filename, 'rt', newline='') as fp:
            reader = csv.DictReader(fp)
            columns = set(reader.fieldnames or [])
            for col in (self.cid_col, self.smiles_col, self.truth_col):
                if col and col not in columns:
                    return (False, f"'{self.filename}': lacks '{col}' column")
        return (True, '')

    def get_dict_reader(self, stack):
        '''
        :param stack: Exit stack.
        :type stack: `contextlib.ExitStack`
        '''
        return csv.DictReader(
            stack.enter_context(
                open_maybe_compressed(self.filename, 'rt', newline='')))

    def get_dict_writer(self, stack, fieldnames=None):
        '''
        :param stack: Exit stack.
        :type stack: `contextlib.ExitStack`

        :param fieldnames: CSV column names.
        :type fieldnames: list(str) or NoneType
        '''

        if fieldnames is None:
            fieldnames = [self.smiles_col, self.cid_col]
            if self.truth_col:
                fieldnames.append(self.truth_col)

        return csv.DictWriter(stack.enter_context(
            open_maybe_compressed(self.filename, 'wt', newline='')),
                              fieldnames=fieldnames)

    @property
    def input_files(self):
        return [self.filename]