model/prediction_result_file.py

"""
Class that represents a file with the prediction results

* allows prediction results to be appended to the file in chunks
* allows results to be skipped if they already exist in the file
"""

import csv
from typing import List, OrderedDict
from pathlib import Path
import logging

from .models.results import PredictionResult
from .errors import PredictionResultExistsError, FileExistsError

LOG = logging.getLogger(__name__)

class PredictionResultsFileIter:
    def __init__(self, results_file_class):
        self._results = results_file_class.get_results()
        self._results_size = len(self._results)
        self._current_index = 0
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self._current_index < self._results_size:
            member = self._results[self._current_index] 
            self._current_index += 1
            return member
        raise StopIteration


class PredictionResultsFile:
    """
    Writes prediction results to a file in chunks
    """
    
    COLNAMES = ['chain_id', 'sequence_md5', 'nres', 'ndom', 'chopping', 'confidence', 'time_sec']

    def __init__(self, 
                 csv_path: Path, *, 
                 chunk_size: int=20, 
                 write_header: bool=True, 
                 allow_append: bool=False, 
                 allow_skip: bool=False):

        self.csv_path = Path(str(csv_path)).absolute()
        self.chunk_size = chunk_size
        self.write_header = write_header
        self.allow_append = allow_append
        self.allow_skip = allow_skip
        self._results_by_id: OrderedDict[PredictionResult] = OrderedDict()
        self._unflushed_results_by_id: OrderedDict[PredictionResult] = OrderedDict()
        self._written_header = False

        self._init()

    def _init(self):
        if self.csv_path.exists():
            msg = f"file '{self.csv_path}' already exists (allow_append={self.allow_append}))"
            LOG.warning(msg)
            if not self.allow_append:
                raise FileExistsError(msg)
            self._read_results()

    def _read_results(self):
        """
        Read the results from the file
        """
        with self.csv_path.open('r') as fp:
            reader = self.get_csv_reader(fp)
            for row in reader:
                if row['chain_id'] == 'chain_id':
                    continue
                pdb_path = Path(row['chain_id'] + '.pdb')
                result = PredictionResult(pdb_path=pdb_path, **row)
                self._results_by_id[result.chain_id] = result

    def get_flushed_results(self):
        self.flush()
        return self._results_by_id.values()

    def get_csv_reader(self, fp):
        return csv.DictReader(fp, delimiter='\t', fieldnames=self.COLNAMES)

    def get_csv_writer(self, fp):
        return csv.DictWriter(fp, delimiter='\t', fieldnames=self.COLNAMES)

    def write_csv_result(self, csv_writer, res: PredictionResult):
        """
        Render PredictionResult result to file pointer
        """
        row = {
            'chain_id': res.chain_id,
            'sequence_md5': res.sequence_md5,
            'nres': res.nres,
            'ndom': res.ndom,
            'chopping': res.chopping if res.chopping is not None else 'NULL',
            'confidence': f'{res.confidence:.3g}' if res.confidence is not None else 'NULL',
            'time_sec': f'{res.time_sec}' if res.time_sec is not None else 'NULL',
        }
        csv_writer.writerow(row)


    def flush(self):
        """
        Write unflushed results to the output file
        """
        if len(self._unflushed_results_by_id) == 0:
            return

        with self.csv_path.open('a') as fp:

            csv_writer = self.get_csv_writer(fp)

            if self.write_header and not self._written_header and len(self._results_by_id) == 0:
                csv_writer.writeheader()
                self._written_header = True

            for result in self._unflushed_results_by_id.values():
                self.write_csv_result(csv_writer, result)

    def has_result(self, result: PredictionResult):
        """
        Check if the prediction result already exists in the file
        """
        return self.has_result_for_chain_id(result.chain_id)

    def has_result_for_chain_id(self, chain_id: str):
        """
        Check if the chain_id already exists in the file
        """
        if chain_id in self._results_by_id:
            return True
        if chain_id in self._unflushed_results_by_id:
            return True

        return False

    def add_result(self, result: PredictionResult):
        """
        Add a result to the buffer (and possibly flush the buffer)
        """
        if self.has_result(result):
            msg = f"result '{result.chain_id}' already exists in file '{self.csv_path}'"
            raise PredictionResultExistsError(msg)

        self._unflushed_results_by_id[result.chain_id] = result
        if len(self._unflushed_results_by_id) >= self.chunk_size:
            self.flush()
            self._results_by_id.update(self._unflushed_results_by_id)
            self._unflushed_results_by_id = OrderedDict()

    def add_results(self, results: List[PredictionResult]):
        for res in results:
            self.add_result(res)

    def __iter__(self):
        return PredictionResultsFileIter(self)