Source code for allele

'''
I don't know how useful this will be, but it allows visualization
on a terminal of the differences between two alleles. I used it
in testing haploh.Haplotyper.
'''
from __future__ import print_function

import const


def compare_allele_files(AF, BF, stream=None, verbose=False):
    '''Visualize two files, *AF* and *BF*, representing alleles, printing dots
    where they are homozygous.
    '''
    with open(AF) as A, open(BF) as B:
        compare_allele_streams(A, B, stream=stream, verbose=verbose)


def compare_allele_streams(A, B, stream=None, verbose=False):
    '''For two input streams, *A* and *B*, visualize their comparison on the output stream,
    printing dots where they are homozygous.
    If *verbose*, print the original streams as well.
    '''
    aline = A.readlines()
    bline = B.readlines()
    for i in range(len(aline)):
        alleleA = aline[i].strip().split(const.SPACE)
        alleleB = bline[i].strip().split(const.SPACE)
        print(visual_compare_alleles(alleleA, alleleB, label='1: '), file=stream)
        print(visual_compare_alleles(alleleB, alleleA, label='2: '), file=stream)
        if verbose:
            print('1: ' + aline[i].strip(), file=stream)
            print('2: ' + bline[i].strip(), file=stream)


def visual_compare_alleles(alleleA, alleleB, label=''):
    '''Return a line with dots for homozygous loci and alleleA otherwise.
    Prefix the line with *label* if given.
    '''
    return label + const.SPACE.join(map(printDot, alleleA, alleleB))


def printDot(alleleA, alleleB):
    '''Print a dot if *alleleA* and *alleleB* are the same, otherwise print *alleleA*
    '''
    if alleleA == alleleB:
        return const.DOT
    else:
        return '%s' % (alleleA)