Mergedb

#!/usr/bin/env python
'''
    The Mergedb program Merges UniProtKB/SwissProt and Uniprot-GOA
    annotation datasets.

    This program takes the following four inputs:
        (1) a UniProtKB/SwissProt file
        (2) a UniProt-GOA file
        (3) a taxon id
        (4) an optional output file

    When the output file name is NOT given, the program will construct
    an output file name by combining the UniProtKB/SwissProt and
    UniProt-GOA file names supplied by the user.

    The GO terms in the UniProtKB/SwissProt file that are NOT in the
    UniProt-GOA file for the supplied taxon id, are merged together
    with those in the UniProt-GOA file and written to the output file.

    How to run this program: 
        For some input UniProtKB/SwissProt file uniprot_sprot.dat.2014_09
    and UniProt-GOA file gene_association.goa_ref_yeast.38 and taxon id
    55929:
        
        python Mergedb -I1=uniprot_sprot.dat.2014_09 -I2=gene_association.goa_ref_yeast.38 -G 559292
       
        One output file will be created: gene_association.goa_ref_yeast.38+sprot.38.1
    which will contain all the entries from the second input file together 
    with all the new annotations for yeast (taxon id 559292) from the first 
    input file.
'''
import os
import sys
from os.path import basename

#import configparser as cp
import ConfigParser as cp
import shutil
import subprocess

import AppendSprot2GOA as as2g
import ArgParser_Mergedb as ap
import Config
import FormatChecker as fc
import LocateDataset as ld

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

# Default configuration file name:
config_filename = '.cafarc' 

class Mergedb:
    def __init__(self):
        # Collect user arguments into a dictionary:
        self.parsed_dict = ap.parse_args()
        # Collect config file entries:
        self.ConfigParam = Config.read_config(config_filename)
        # Extract workspace name
        self.work_dir = self.ConfigParam['workdir']
        # Create workspace, if one does not exist:
        if not os.path.exists(self.work_dir):
            os.makedirs(self.work_dir)
        # Extract filename at time t1:
        t1 = self.parsed_dict['t1']
        # Extract filename at time t2:
        t2 = self.parsed_dict['t2']
        # Extract output filename:
        outfile_basename = basename(self.parsed_dict['outfile'])
        # Locate t1 file:
        self.t1_input_file = ld.locate_SwissProtfile(t1,
                                                     self.work_dir)
        # Locate t2 file:
        self.t2_input_file = ld.locate_GOAfile(t2,
                                               self.work_dir)
        # Create output file name together with file version number:
        self.output_filename = self.create_outfilename(self.parsed_dict,
                                                  outfile_basename,
                                                  self.work_dir)
        return None

    def create_iterator(self, infile):
        # Returns an iterator object for an input uniprot-goa file along 
        # with a list of all field names contained in the uniprot-goa file
        infile_handle = open(infile, 'r')
        iter_handle = GOA.gafiterator(infile_handle)
        for ingen in iter_handle:
            if len(ingen) == 17:
                GAFFIELDS = GOA.GAF20FIELDS
                break
            else:
                GAFFIELDS = GOA.GAF10FIELDS
                break
        infile_handle = open(infile, 'r')
        iter_handle = GOA.gafiterator(infile_handle)
        return iter_handle, GAFFIELDS

    def create_outfilename(self, params, outfile, work_dir):
        """
         This method creates an output filename based on the output 
         file prefix provided by the user and returns the created 
         file name. 
        """
        if not outfile == '':
            ob = basename(self.parsed_dict['outfile'])
        else:
            ob = basename(self.parsed_dict['t2']) + '+sprot.' + \
                          str(basename(self.parsed_dict['t1']).split('.')[-1])
        index = 1
        while os.path.exists(self.work_dir + '/' + ob + '.' + str(index)):
            index = index + 1
        output_filename = self.work_dir + '/' + ob + '.' + str(index)
        return output_filename

    def check_sprot_format(self, sprot_fname):
        """
        This method exits the Mergedb program on any of the
        following conditions:
            Case 1: if the file is empty
            Case 2: if the file is NOT in UniProtKB/SwissProt format.
                    To check this it invokes check_sprot_format method
                    of FormatChecker module.
        """
        if os.stat(sprot_fname).st_size == 0:
            print bcolors.WARNING + 'You submitted an empty file: ' + sprot_fname + \
                  bcolors.ENDC
            sys.exit(1)
        elif not fc.check_sprot_format(open(sprot_fname, 'r')):
            print bcolors.WARNING + 'File format error: ' + \
                  basename(sprot_fname) + bcolors.ENDC
            print bcolors.WARNING + 'File must be in UniProtKB/SwissProt ' + \
                'format' + bcolors.ENDC
            sys.exit(1)

    def check_gaf_format(self, goa_fname):
        """
        This method exits the Benchmark program on any of the
        following conditions:
            Case 1: if the file is empty
            Case 2: if the file is NOT in GAF format. To check this
                    it invokes check_gaf_format method of
                    FormatChecker module.
        """
        if os.stat(goa_fname).st_size == 0:
            print bcolors.WARNING + 'You submitted an empty file: ' + goa_fname + \
                  bcolors.ENDC
            sys.exit(1)
        elif not fc.check_gaf_format(open(goa_fname, 'r')):
            print bcolors.WARNING + "File format error: " + \
                  basename(goa_fname) + bcolors.ENDC
            print bcolors.WARNING + 'File must be in GAF 1.0 or GAF 2.0 ' + \
                'format' + bcolors.ENDC
            sys.exit(1)

    def print_prolog(self):
        print("*************************************************")
        print("Running Merge Database Tool !!!!!")
        print('Following is a list of user supplied inputs:')
        for arg in self.parsed_dict:
            print(arg + ': ' + str(self.parsed_dict[arg]))
        print('*********************************************\n')
        return None

    def print_epilog(self, goCount):
        if os.path.exists(self.output_filename):
            print(bcolors.OKGREEN + 'The following output file is created:' + \
                  bcolors.ENDC)
            print('    ' + basename(self.output_filename))
            print('    ' + str(goCount) + ' protein annotations added')
        else:
            print(bcolors.WARNING + 'No output file is created with the ' + \
                  'given input parameters' + bcolors.ENDC)
        print(bcolors.OKGREEN + 'Thank you for using Merge Database Tool' + \
              bcolors.ENDC)
        return None

    def process_data(self): 
        # Print wellcome message:
        self.print_prolog()

        # Check UniProtKB/SwissProt file format: 
        self.check_sprot_format(self.t1_input_file)

        # Check UniProt-GOA file format:
        self.check_gaf_format(self.t2_input_file) 

        # Merging in TWO steps:
        print ('Merging records in two steps - copying and appending:')

        # Step 1: 
           # Copy the gene association file to the output file:
        print ('Copying records from ' + \
                basename(self.t2_input_file) + ' to ' + \
                basename(self.output_filename) + ' ...')

        cmd = "cp " + self.t2_input_file + " " + self.output_filename
        subprocess.call(cmd, shell=True)

        # Step 2:
            # Fetch records from Uniprot-SwissProt file
            # Check for duplicacy in UniProt-GOA file 
            # Convert them to GOA records
            # Append them at the end of the output file one by one
            # All these are performed in appendSprot2goa method
        print ('Appedning records from ' + \
                basename(self.t1_input_file) + ' to ' + \
                basename(self.output_filename) + ' ...')

        goCount = as2g.appendSprot2goa(open(self.t1_input_file, 'r'),
                                            self.t2_input_file,
                                            self.parsed_dict['g'],
                                       open(self.output_filename, 'a'))

        # Print the summary of running this program:
        self.print_epilog(goCount)
        return None

if __name__ == '__main__':
    if len(sys.argv) == 1:
        print (sys.argv[0] + ':')
        print(__doc__)
    else: 
        # Creates an instance of Mergedb class:
        md = Mergedb()    
        # Process data and create merged file:
        md.process_data() 
    sys.exit(0)