<commit_before>#!/usr/bin/env python
# encoding: utf-8

import os
import re
import sys
import gzip
import pysam
import itertools
import numpy as np
from fstats import fstats
from collections import OrderedDict
from itertools import combinations, izip_longest

def set_header(vcf_path):
    """Open VCF file and read in header line as Ordered Dict"""

    vcf_file = gzip.open(vcf_path,'rb')
    header_dict = None
    for line in vcf_file:
        if line.startswith("#CHROM"):
            header = line.strip("#").strip().split()
            header_dict = OrderedDict([(item, None) for item in header])
            break

    vcf_file.close()
    return header_dict

def parse_populations_list(populations):
    populations_dict  = {}
    for pop in populations:
        pop_name, sample_ids = pop.strip().split(":")
        sample_ids = sample_ids.split(",")
        populations_dict[pop_name] = sample_ids

    return populations_dict

def pairwise(iterable):
    """Generates paris of slices from iterator
       s -> (s0,s1), (s1,s2), (s2, s3), ..."""

    a, b = itertools.tee(iterable)
    next(b, None)
    return itertools.izip(a, b)

def get_slice_indicies(vcf_bgzipped_file, window_size=500):
    """Get slice information from VCF file that is tabix indexed file (bgzipped). """

    # READ IN FILE HEADER
    tbx = pysam.Tabixfile(vcf_bgzipped_file) # TODO: create try statement to test that file is actually a VCF
    chrms = tbx.contigs

    # PARSE LENGTH INFO FROM HEADER
    chrm_lengths = []
    for line in tbx.header:

        if line.startswith("##contig="):

            chrm_name = re.findall(r'ID=.*,',line)
            chrm_name = chrm_name[0].strip('ID=').strip(',')

            chrm_length = re.findall(r'length=.*>',line)
            chrm_length = int(chrm_length[0].strip('length=').strip('>'))

            chrm_lengths.append((chrm_name, 1, chrm_length))

    chrm_lengths = tuple(chrm_lengths)
    tbx.close()

    # GENERATE SLICES
    # current does not make overlapping slices.

    if regions == None:

        for chrm, start, stop in chrm_lengths:

            slice_indicies = itertools.islice(xrange(start, stop + 1), 0, stop + 1, window_size)

            for count, s in enumerate(pairwise(slice_indicies)):
                yield (chrm, s[0], s[1] -1) # subtract one to prevent 1 bp overlap
    else:

        chrm, start, stop = re.split(r':|-', regions)
        start, stop = int(start), int(stop)

        slice_indicies = itertools.islice(xrange(start, stop + 1), 0, stop + 1, window_size)
        for count, s in enumerate(pairwise(slice_indicies)):
            yield (chrm, s[0], s[1] -1) # subtract one to prevent 1 bp overlap

def slice_vcf(vcf_bgzipped_file, chrm, start, stop):

    tbx = pysam.Tabixfile(vcf_bgzipped_file)
    return tuple(row for row in tbx.fetch(chrm, start, stop))

def parse_info_field(info_field):

    info_dict = {}
    for item in info_field.split(';'):
        pair = item.split("=")
        if len(pair) == 2:
            info_dict[pair[0]] = pair[1] # this could be improved on
    return info_dict

def parse_vcf_line(pos, header):
    """Read in VCF line and convert it to an OrderedDict"""

    pos_parts = pos.strip().split()
    vcf_line_dict = header

    for count, item in enumerate(vcf_line_dict):
        vcf_line_dict[item] = pos_parts[count]

    sample_format = vcf_line_dict["FORMAT"].split(":")

    for count, item in enumerate(vcf_line_dict):

        if count >= 9:
            genotype = vcf_line_dict[item]

            if genotype == "./." or genotype == ".":      # "'./.'' for dip, '.' for haploid
                vcf_line_dict[item] = None

            else:
                genotype = dict(zip(sample_format,genotype.split(":")))

                # CONVERT STRINGS TO APPOPRIATE TYPES (INTS, FLOATS, ETC.)
                genotype['GQ'] = float(genotype['GQ'])
                genotype['DP'] = int(genotype['DP'])
                genotype['AD'] = tuple(int(ad) for ad in genotype['AD'].split(","))
                genotype['PL'] = tuple(int(ad) for ad in genotype['PL'].split(","))

                vcf_line_dict[item] = genotype

    vcf_line_dict['POS'] = int(vcf_line_dict['POS'])
    vcf_line_dict['QUAL'] = float(vcf_line_dict['QUAL'])

    vcf_line_dict['INFO'] = parse_info_field(vcf_line_dict['INFO'])

    return vcf_line_dict

def filter_on_population_sizes(vcfline, populations, min_samples=0):

    sample_counts = []

    for pop in populations.keys():
        sample_count = 0
        for sample in populations[pop]:
            if vcfline[sample] != None:
                sample_count += 1

        sample_counts.append(sample_count)

    sample_counts = [item for item in sample_counts if item >= min_samples]

    if len(sample_counts) >= 2:
        return True
    else:
        return False

def format_output(chrm, start, stop, depth, stat_id, multilocus_f_statistics):

    line = ",".join((chrm, str(start), str(stop), str(depth)))
    for pair in multilocus_f_statistics.keys():
        if multilocus_f_statistics[pair] == None:
            continue

        if multilocus_f_statistics[pair][stat_id] == np.nan:
            value = 'NA'

        else:
            value = str(multilocus_f_statistics[pair][stat_id])

        line += "," + value

    line += "\n"
    return (line)

def calc_allele_counts(populations, vcf_line_dict):

    allele_counts = populations.fromkeys(populations.keys(), None)

    for population in populations.keys():

        allele_format_dict = {0:0.0,1:0.0,2:0.0,3:0.0,4:0.0}   # create dict to prevent pointer issues
        allele_counts[population] = allele_format_dict

        for sample_id in populations[population]:

            if vcf_line_dict[sample_id] != None:

                genotype = vcf_line_dict[sample_id]
                genotype = genotype["GT"].split("/")

                if genotype == [".","."]: continue

                genotype = [int(item) for item in genotype]

                for allele in genotype:
                    allele_counts[population][allele] += 1.0

    return allele_counts

def calc_fstats(populations, allele_counts):

    #test_pair = {'So_Riviere_Goyaves': np.array([ 0.0, 1.0, 0.0, 0.0]), 'Plage_de_Viard': np.array([ 1.0, 0.0, 0.0, 0.0]),}

    # CALCULATE ALLELE FREQUENCIES
    allele_freqs_dict = populations.fromkeys(populations.keys(), None)
    for population in allele_counts.keys():
        counts =  allele_counts[population].values()

        if sum(counts) == 0.0:
            freqs = [0.0] * 4

        else:
            freqs =  [count/sum(counts) for count in counts]

        allele_freqs_dict[population] = freqs

    allele_freqs = allele_freqs_dict.values()

    # CACULATE PAIRWISE F-STATISTICS
    pairwise_results = {}
    for population_pair in combinations(populations.keys(),2):

        pop1, pop2 =  population_pair
        Ns = [sum(allele_counts[pop].values()) for pop in [pop1, pop2]]

        if 0 in Ns:
            values = [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]
            values_dict = dict(zip(['Hs_est', 'Ht_est', 'Gst_est',
                                    'G_prime_st_est', 'G_double_prime_st_est',
                                    'D_est'],
                                    values))
            pairwise_results[population_pair] = values_dict
            continue

        else:

            pop1 = allele_freqs_dict[pop1]
            pop2 = allele_freqs_dict[pop2]

            allele_freqs = [pop1, pop2]

            n = float(len(Ns))
            Ns_harm = fstats.harmonic_mean(Ns)
            Ns_harm_chao = fstats.harmonic_mean_chao(Ns)

            # CALCULATE Hs AND Ht
            n = 2
            Hs_prime_est_ = fstats.Hs_prime_est(allele_freqs, n)
            Ht_prime_est_ = fstats.Ht_prime_est(allele_freqs, n)
            Hs_est_ = fstats.Hs_est(Hs_prime_est_, Ns_harm)
            Ht_est_ = fstats.Ht_est(Ht_prime_est_, Hs_est_, Ns_harm, n)

            # CALCULATE F-STATISTICS
            Gst_est_ = fstats.Gst_est(Ht_est_, Hs_est_)
            G_prime_st_est_ = fstats.G_prime_st_est(Ht_est_, Hs_est_, Gst_est_, n)
            G_double_prime_st_est_ = fstats.G_double_prime_st_est(Ht_est_, Hs_est_, n)
            D_est_ = fstats.D_est(Ht_est_, Hs_est_, n)

            # PRINT OUTPUT
            values = [Hs_est_, Ht_est_, Gst_est_, G_prime_st_est_,
                      G_double_prime_st_est_, D_est_]
            values_dict = dict(zip(['Hs_est', 'Ht_est', 'Gst_est', \
                                    'G_prime_st_est', 'G_double_prime_st_est',
                                    'D_est'],
                                    values))

            pairwise_results[population_pair] = values_dict

    return pairwise_results

def output_header():
    return 'Chrm,Start,Stop,DP_mean,DP_std'

def calculate_multilocus_f_statistics(Hs_est_dict, Ht_est_dict):

    multilocus_f_statistics = {}

    for key in Hs_est_dict.keys():

        Hs_est_list = Hs_est_dict[key]
        Ht_est_list = Ht_est_dict[key]

        pairs = zip(Hs_est_list, Ht_est_list)
        pairs = [pair for pair in pairs if np.nan not in pair]

        multilocus_f_statistics[key] = None

        if len(pairs) != 0:

            n = 2 # fix this
            Gst_est = fstats.multilocus_Gst_est(Ht_est_list, Hs_est_list)
            G_prime_st_est = fstats.multilocus_G_prime_st_est(Ht_est_list, Hs_est_list, n)
            G_double_prime_st_est = fstats.multilocus_G_double_prime_st_est(Ht_est_list, Hs_est_list, n)
            D_est = fstats.multilocus_D_est(Ht_est_list, Hs_est_list, n)

            values_dict = dict(zip(['Gst_est', 'G_prime_st_est', 'G_double_prime_st_est', 'D_est'],\
                          [Gst_est, G_prime_st_est, G_double_prime_st_est, D_est]))

            multilocus_f_statistics[key] = values_dict

    return multilocus_f_statistics

def update_Hs_and_Ht_dicts(f_statistics, Hs_est_dict, Ht_est_dict):

    for pop_pair in f_statistics.keys():

        if Hs_est_dict.has_key(pop_pair) == True:
            Hs_est_dict[pop_pair].append(f_statistics[pop_pair]['Hs_est'])
        else:
            Hs_est_dict[pop_pair] = [f_statistics[pop_pair]['Hs_est']]

        if Ht_est_dict.has_key(pop_pair) == True:
            Ht_est_dict[pop_pair].append(f_statistics[pop_pair]['Ht_est'])
        else:
            Ht_est_dict[pop_pair] = [f_statistics[pop_pair]['Ht_est']]
    return (Hs_est_dict, Ht_est_dict)

def multilocus_f_statistics_2_sorted_list(multilocus_f_statistics, order):

    order = []
    if len(order) == 0:
        for pair in multilocus_f_statistics.keys():
            joined_pair = '.'.join(pair)
            for stat in multilocus_f_statistics[pair].keys():
                order.append('.'.join((joined_pair,stat)))
    order.sort()
    stats = []
    for key in order:
        pop1, pop2, stat = key.split(".")
        stat = multilocus_f_statistics[(pop1,pop2)][stat]
        stats.append(stat)

    return (stats, order)

def calc_slice_stats(data):
    """Main function for caculating statistics.

       Make it easy to add more statistics.
    """

    tabix_slice, chrm, start, stop, populations, header = data

    if len(tabix_slice) == 0:
        return None

    else:
        # Create lists to store values for final printing.
        output_line = [chrm, start, stop]
        total_depth = []
        snp_wise_f_statistics = []

        Hs_est_dict = {}
        Ht_est_dict = {}
        ordered_list = None
        snp_count = 0

        for count, line in enumerate(tabix_slice):

            vcf_line_dict = parse_vcf_line(line, header)

            # CREATE FILTERS HERE:
            if filter_on_population_sizes(vcf_line_dict, populations, min_samples=5) == False:
                continue

            # CALCULATE SNPWISE F-STATISTICS
            allele_counts = calc_allele_counts(populations, vcf_line_dict)
            f_statistics = calc_fstats(populations, allele_counts)

            # UPDATE Hs AND Ht DICTIIONARIES
            Hs_est_dict, Ht_est_dict = update_Hs_and_Ht_dicts(f_statistics, Hs_est_dict, Ht_est_dict)

            f_statistics['LOCATION'] = (chrm, start, stop)
            snp_wise_f_statistics.append(f_statistics)

            total_depth.append(int(vcf_line_dict['INFO']["DP"]))
            snp_count = count

        multilocus_f_statistics = calculate_multilocus_f_statistics(Hs_est_dict, Ht_est_dict)

        # UPDATE OUTPUT LINE WITH DEPTH INFO
        total_depth = np.array(total_depth)
        output_line += [snp_count, total_depth.mean(), total_depth.std()]

        return ([chrm, start, stop, snp_count, total_depth.mean(), total_depth.std()], multilocus_f_statistics)

<commit_msg>Minor cleanups<commit_after>#!/usr/bin/env python
# encoding: utf-8

import os
import re
import sys
import gzip
import pysam
import itertools
import numpy as np
from fstats import fstats
from collections import OrderedDict
from itertools import combinations, izip_longest

def set_header(vcf_path):
    """Open VCF file and read in header line as Ordered Dict"""

    vcf_file = gzip.open(vcf_path,'rb')
    header_dict = None
    for line in vcf_file:
        if line.startswith("#CHROM"):
            header = line.strip("#").strip().split()
            header_dict = OrderedDict([(item, None) for item in header])
            break

    vcf_file.close()
    return header_dict

def parse_populations_list(populations):
    populations_dict  = {}
    for pop in populations:
        pop_name, sample_ids = pop.strip().split(":")
        sample_ids = sample_ids.split(",")
        populations_dict[pop_name] = sample_ids

    return populations_dict

def pairwise(iterable):
    """Generates paris of slices from iterator
       s -> (s0,s1), (s1,s2), (s2, s3), ..."""

    a, b = itertools.tee(iterable)
    next(b, None)
    return itertools.izip(a, b)

def get_slice_indicies(vcf_bgzipped_file, regions, window_size):
    """Get slice information from VCF file that is tabix indexed file (bgzipped). """

    # READ IN FILE HEADER
    tbx = pysam.Tabixfile(vcf_bgzipped_file) # TODO: create try statement to test that file is actually a VCF
    chrms = tbx.contigs

    # PARSE LENGTH INFO FROM HEADER
    chrm_lengths = []
    for line in tbx.header:

        if line.startswith("##contig="):

            chrm_name = re.findall(r'ID=.*,',line)
            chrm_name = chrm_name[0].strip('ID=').strip(',')

            chrm_length = re.findall(r'length=.*>',line)
            chrm_length = int(chrm_length[0].strip('length=').strip('>'))

            chrm_lengths.append((chrm_name, 1, chrm_length))

    chrm_lengths = tuple(chrm_lengths)
    tbx.close()

    # GENERATE SLICES
    # current does not make overlapping slices.

    if regions == None:

        for chrm, start, stop in chrm_lengths:

            slice_indicies = itertools.islice(xrange(start, stop + 1), 0, stop + 1, window_size)

            for count, s in enumerate(pairwise(slice_indicies)):
                yield (chrm, s[0], s[1] -1) # subtract one to prevent 1 bp overlap
    else:

        chrm, start, stop = re.split(r':|-', regions)
        start, stop = int(start), int(stop)

        slice_indicies = itertools.islice(xrange(start, stop + 1), 0, stop + 1, window_size)
        for count, s in enumerate(pairwise(slice_indicies)):
            yield (chrm, s[0], s[1] -1) # subtract one to prevent 1 bp overlap

def slice_vcf(vcf_bgzipped_file, chrm, start, stop):

    tbx = pysam.Tabixfile(vcf_bgzipped_file)
    return tuple(row for row in tbx.fetch(chrm, start, stop))

def parse_info_field(info_field):

    info_dict = {}
    for item in info_field.split(';'):
        pair = item.split("=")
        if len(pair) == 2:
            info_dict[pair[0]] = pair[1] # this could be improved on
    return info_dict

def parse_vcf_line(pos, header):
    """Read in VCF line and convert it to an OrderedDict"""

    pos_parts = pos.strip().split()
    vcf_line_dict = header

    for count, item in enumerate(vcf_line_dict):
        vcf_line_dict[item] = pos_parts[count]

    sample_format = vcf_line_dict["FORMAT"].split(":")

    for count, item in enumerate(vcf_line_dict):

        if count >= 9:
            genotype = vcf_line_dict[item]

            if genotype == "./." or genotype == ".":      # "'./.'' for dip, '.' for haploid
                vcf_line_dict[item] = None

            else:
                genotype = dict(zip(sample_format,genotype.split(":")))

                # CONVERT STRINGS TO APPOPRIATE TYPES (INTS, FLOATS, ETC.)
                genotype['GQ'] = float(genotype['GQ'])
                genotype['DP'] = int(genotype['DP'])
                genotype['AD'] = tuple(int(ad) for ad in genotype['AD'].split(","))
                genotype['PL'] = tuple(int(ad) for ad in genotype['PL'].split(","))

                vcf_line_dict[item] = genotype

    vcf_line_dict['POS'] = int(vcf_line_dict['POS'])
    vcf_line_dict['QUAL'] = float(vcf_line_dict['QUAL'])

    vcf_line_dict['INFO'] = parse_info_field(vcf_line_dict['INFO'])

    return vcf_line_dict

def filter_on_population_sizes(vcfline, populations, min_samples=0):

    sample_counts = []

    for pop in populations.keys():
        sample_count = 0
        for sample in populations[pop]:
            if vcfline[sample] != None:
                sample_count += 1

        sample_counts.append(sample_count)

    sample_counts = [item for item in sample_counts if item >= min_samples]

    if len(sample_counts) >= 2:
        return True
    else:
        return False

def format_output(chrm, start, stop, depth, stat_id, multilocus_f_statistics):

    line = ",".join((chrm, str(start), str(stop), str(depth)))
    for pair in multilocus_f_statistics.keys():
        if multilocus_f_statistics[pair] == None:
            continue

        if multilocus_f_statistics[pair][stat_id] == np.nan:
            value = 'NA'

        else:
            value = str(multilocus_f_statistics[pair][stat_id])

        line += "," + value

    line += "\n"
    return (line)

def calc_allele_counts(populations, vcf_line_dict):

    allele_counts = populations.fromkeys(populations.keys(), None)

    for population in populations.keys():

        allele_format_dict = {0:0.0,1:0.0,2:0.0,3:0.0,4:0.0}   # create dict to prevent pointer issues
        allele_counts[population] = allele_format_dict

        for sample_id in populations[population]:

            if vcf_line_dict[sample_id] != None:

                genotype = vcf_line_dict[sample_id]
                genotype = genotype["GT"].split("/")

                if genotype == [".","."]: continue

                genotype = [int(item) for item in genotype]

                for allele in genotype:
                    allele_counts[population][allele] += 1.0

    return allele_counts

def calc_fstats(populations, allele_counts):

    # CALCULATE ALLELE FREQUENCIES
    allele_freqs_dict = populations.fromkeys(populations.keys(), None)
    for population in allele_counts.keys():
        counts =  allele_counts[population].values()

        if sum(counts) == 0.0:
            freqs = [0.0] * 4

        else:
            freqs =  [count/sum(counts) for count in counts]

        allele_freqs_dict[population] = freqs

    allele_freqs = allele_freqs_dict.values()

    # CACULATE PAIRWISE F-STATISTICS
    pairwise_results = {}
    for population_pair in combinations(populations.keys(),2):

        pop1, pop2 =  population_pair
        Ns = [sum(allele_counts[pop].values()) for pop in [pop1, pop2]]

        if 0 in Ns:
            values = [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]
            values_dict = dict(zip(['Hs_est', 'Ht_est', 'Gst_est',
                                    'G_prime_st_est', 'G_double_prime_st_est',
                                    'D_est'],
                                    values))
            pairwise_results[population_pair] = values_dict
            continue

        else:

            pop1 = allele_freqs_dict[pop1]
            pop2 = allele_freqs_dict[pop2]

            allele_freqs = [pop1, pop2]

            n = float(len(Ns))
            Ns_harm = fstats.harmonic_mean(Ns)
            Ns_harm_chao = fstats.harmonic_mean_chao(Ns)

            # CALCULATE Hs AND Ht
            n = 2
            Hs_prime_est_ = fstats.Hs_prime_est(allele_freqs, n)
            Ht_prime_est_ = fstats.Ht_prime_est(allele_freqs, n)
            Hs_est_ = fstats.Hs_est(Hs_prime_est_, Ns_harm)
            Ht_est_ = fstats.Ht_est(Ht_prime_est_, Hs_est_, Ns_harm, n)

            # CALCULATE F-STATISTICS
            Gst_est_ = fstats.Gst_est(Ht_est_, Hs_est_)
            G_prime_st_est_ = fstats.G_prime_st_est(Ht_est_, Hs_est_, Gst_est_, n)
            G_double_prime_st_est_ = fstats.G_double_prime_st_est(Ht_est_, Hs_est_, n)
            D_est_ = fstats.D_est(Ht_est_, Hs_est_, n)

            # PRINT OUTPUT
            values = [Hs_est_, Ht_est_, Gst_est_, G_prime_st_est_,
                      G_double_prime_st_est_, D_est_]
            values_dict = dict(zip(['Hs_est', 'Ht_est', 'Gst_est', \
                                    'G_prime_st_est', 'G_double_prime_st_est',
                                    'D_est'],
                                    values))

            pairwise_results[population_pair] = values_dict

    return pairwise_results

def output_header():
    return 'Chrm,Start,Stop,DP_mean,DP_std'

def calculate_multilocus_f_statistics(Hs_est_dict, Ht_est_dict):

    multilocus_f_statistics = {}

    for key in Hs_est_dict.keys():

        Hs_est_list = Hs_est_dict[key]
        Ht_est_list = Ht_est_dict[key]

        pairs = zip(Hs_est_list, Ht_est_list)
        pairs = [pair for pair in pairs if np.nan not in pair]

        multilocus_f_statistics[key] = None

        if len(pairs) != 0:

            n = 2 # fix this
            Gst_est = fstats.multilocus_Gst_est(Ht_est_list, Hs_est_list)
            G_prime_st_est = fstats.multilocus_G_prime_st_est(Ht_est_list, Hs_est_list, n)
            G_double_prime_st_est = fstats.multilocus_G_double_prime_st_est(Ht_est_list, Hs_est_list, n)
            D_est = fstats.multilocus_D_est(Ht_est_list, Hs_est_list, n)

            values_dict = dict(zip(['Gst_est', 'G_prime_st_est', 'G_double_prime_st_est', 'D_est'],\
                          [Gst_est, G_prime_st_est, G_double_prime_st_est, D_est]))

            multilocus_f_statistics[key] = values_dict

    return multilocus_f_statistics

def update_Hs_and_Ht_dicts(f_statistics, Hs_est_dict, Ht_est_dict):

    for pop_pair in f_statistics.keys():

        if Hs_est_dict.has_key(pop_pair) == True:
            Hs_est_dict[pop_pair].append(f_statistics[pop_pair]['Hs_est'])
        else:
            Hs_est_dict[pop_pair] = [f_statistics[pop_pair]['Hs_est']]

        if Ht_est_dict.has_key(pop_pair) == True:
            Ht_est_dict[pop_pair].append(f_statistics[pop_pair]['Ht_est'])
        else:
            Ht_est_dict[pop_pair] = [f_statistics[pop_pair]['Ht_est']]
    return (Hs_est_dict, Ht_est_dict)

def multilocus_f_statistics_2_sorted_list(multilocus_f_statistics, order):

    order = []
    if len(order) == 0:
        for pair in multilocus_f_statistics.keys():
            joined_pair = '.'.join(pair)
            for stat in multilocus_f_statistics[pair].keys():
                order.append('.'.join((joined_pair,stat)))
    order.sort()
    stats = []
    for key in order:
        pop1, pop2, stat = key.split(".")
        stat = multilocus_f_statistics[(pop1,pop2)][stat]
        stats.append(stat)

    return (stats, order)

def calc_slice_stats(data):
    """Main function for caculating statistics.

       Make it easy to add more statistics.
    """

    tabix_slice, chrm, start, stop, populations, header = data

    if len(tabix_slice) == 0:
        return None

    else:
        # Create lists to store values for final printing.
        output_line = [chrm, start, stop]
        total_depth = []
        snp_wise_f_statistics = []

        Hs_est_dict = {}
        Ht_est_dict = {}
        ordered_list = None
        snp_count = 0

        for count, line in enumerate(tabix_slice):

            vcf_line_dict = parse_vcf_line(line, header)

            # CREATE FILTERS HERE:
            if filter_on_population_sizes(vcf_line_dict, populations, min_samples=5) == False:
                continue

            # CALCULATE SNPWISE F-STATISTICS
            allele_counts = calc_allele_counts(populations, vcf_line_dict)
            f_statistics = calc_fstats(populations, allele_counts)

            # UPDATE Hs AND Ht DICTIIONARIES
            Hs_est_dict, Ht_est_dict = update_Hs_and_Ht_dicts(f_statistics, Hs_est_dict, Ht_est_dict)

            f_statistics['LOCATION'] = (chrm, start, stop)
            snp_wise_f_statistics.append(f_statistics)

            total_depth.append(int(vcf_line_dict['INFO']["DP"]))
            snp_count = count

        multilocus_f_statistics = calculate_multilocus_f_statistics(Hs_est_dict, Ht_est_dict)

        # UPDATE OUTPUT LINE WITH DEPTH INFO
        total_depth = np.array(total_depth)
        output_line += [snp_count, total_depth.mean(), total_depth.std()]

        return ([chrm, start, stop, snp_count, total_depth.mean(), total_depth.std()], multilocus_f_statistics)

<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8

import re
import sys
import gzip
import pysam
#import itertools
import mimetypes
from collections import OrderedDict  #,defaultdict

class VCF(object):
    """docstring for VCF"""

    def __init__(self,
                 input,
                 output=None,
                 populations=None,
                 region=None,
                 window_size=1,
                 step=0,
                 snvs=None):

        super(VCF, self).__init__()

        self.input = input
        self.output = output
        self.step = step
        self.populations = populations
        self.region = region
        self.window_size = window_size
        self.header = self.__extract_header__()
        self.chrms_2_sizes = self.__get_chrm_ids_and_sizes__()
        self.empty_vcf_line = self.make_empty_vcf_ordered_dict()

    def __open_vcf__(self):
        """Open vcf file as gzip or as text."""

        if type(self.input) is file:
            fin = self.input

        elif mimetypes.guess_type(self.input)[-1] is 'gzip':
            fin = gzip.open(self.input, 'rb')

        else:
            fin = open(self.input, 'r')

        return fin

    def __extract_header__(self):
        #print [line.strip() for line in self.__open_vcf__()]

        header = []
        for line in self.__open_vcf__():

            if line.startswith("#") == True:
                header.append(line.strip())
            else:
                break

        return header

    def __get_chrm_ids_and_sizes__(self):
        """ Extract chromosome ids and sizes from vcf file.
            Return as dictionary"""

        chrms_sizes_dict = OrderedDict()
        # with self.__open_vcf__() as fin:

        for line in self.header:

            if line.startswith("##contig"):
                chrm_name = re.findall(r'ID=.*,', line)
                chrm_name = chrm_name[0].strip('ID=').strip(',')

                chrm_length = re.findall(r'length=.*>', line)
                chrm_length = int(chrm_length[0].strip('length=').strip('>'))

                chrms_sizes_dict[chrm_name] = chrm_length
                #break

        return chrms_sizes_dict

    def make_empty_vcf_ordered_dict(self):
        """Open VCF file and read in #CHROM line as an Ordered Dict"""

        header_dict = None
        for line in self.header:
            if line.startswith("#CHROM"):
                header = line.strip("#").strip().split()
                header_dict = OrderedDict([(item, None) for item in header])
                break

        return header_dict

    def process_snp(self, snp_call):
        if snp_call == "0/0":
            return (0,0)

        elif snp_call == "1/1":
            return (1,1)

        elif snp_call == '1/0' or \
               snp_call == '0/1':
            return (0,1)

        # skip multiallelic sites
        else:
            return None

    def process_snp_call(self, snp_call, ref, alt, IUPAC_ambiguities=False):
        """Process VCF genotype fields."""

        # IUPAC ambiguity codes
        IUPAC_dict = {('A', 'C'): 'M',
                      ('A', 'G'): 'R',
                      ('A', 'T'): 'W',
                      ('C', 'G'): 'S',
                      ('C', 'T'): 'Y',
                      ('G', 'T'): 'K',
                      ('A', 'C', 'G'): 'V',
                      ('A', 'C', 'T'): 'H',
                      ('A', 'G', 'T'): 'D',
                      ('C', 'G', 'T'): 'B'}

        if snp_call == None:
            called_base = "-"

        if snp_call["GT"] == "./.":
            called_base = "-"

        else:
            allele1, allele2 = snp_call["GT"].split("/")

            # process "0/0"
            if allele1 == '0' and allele2 == '0':
                called_base = ref

            if allele1 == '1' and allele2 == '1':
                called_base = alt

            # process "0/N"
            if allele1 == '0' and allele2 != '0':

                if IUPAC_ambiguities == False:
                    called_base = 'N'

                else:
                    call = [ref] + [alt.split(',')[int(allele2) - 1]]
                    call.sort()
                    call = tuple(call)
                    called_base = IUPAC_dict[call]

            # process "2/2, 1/2, etc."
            if int(allele1) >= 1 and int(allele2) > 1:

                # deal with homozygotes
                if allele1 == allele2:
                    called_base = alt.split(',')[int(allele1) - 1]

                # deal with heterozygotes
                else:

                    if IUPAC_ambiguities == False:
                        called_base = 'N'

                    else:
                        ref = alt.split(',')[int(allele1) - 1]
                        alt = alt.split(',')[int(allele2) - 1]
                        call = [ref, alt]
                        call.sort()
                        call = tuple(call)
                        called_base = IUPAC_dict[call]

        return called_base

    def count_alleles(self, chunk):

        results = []
        for line in chunk:
            pop_counts = {}

            for pop in self.populations.keys():
                allele_counts = {'REF':0, 'ALT':0}

                for sample in self.populations[pop]:

                    if line[sample] != None:
                        ref, alt = self.process_snp(line[sample]['GT'])
                        allele_counts['REF'] += ref
                        allele_counts['ALT'] += alt

                pop_counts[pop] = allele_counts.copy()

            results.append(pop_counts.copy())

        return results

    def vcf_slice_iterator(self, vcf_bgzipped_file, region):

        tbx = pysam.Tabixfile(vcf_bgzipped_file)

        try:
            vcf_slice = tbx.fetch(*region)
        except ValueError:
            print 'bad vcf slice:', region
            sys.exit()
        else:
            for row in vcf_slice:
                yield self.parse_vcf_line(row, self.empty_vcf_line)
            #return ((chrm, start, stop), chunk)
            #return tuple(row for row in vcf_slice)

    def vcf_file_iterator(self, as_dict=True):

        for line in self.__open_vcf__():
            if line.startswith("#") is not True:
                if as_dict == True:
                    yield self.parse_vcf_line(line, self.empty_vcf_line.copy())
                else:
                    yield line
            else:
                continue

    def parse_info_field(self, info_field):

        info_dict = {}
        for item in info_field.split(';'):
            pair = item.split("=")
            if len(pair) == 2:
                info_dict[pair[0]] = pair[1]   # this could be improved on

        return info_dict

    def get_population_sizes(self, vcfline, populations):

        sample_counts = {}

        for pop in populations.keys():
            sample_count = 0

            for sample in populations[pop]:
                if vcfline[sample] is not None:
                    sample_count += 1

            sample_counts[pop] = sample_count

        return sample_counts

    def parse_vcf_line(self, pos, vcf_line_dict):
        """Read in VCF line and convert it to an OrderedDict"""

        pos_parts = pos.strip().split()
        for count, item in enumerate(vcf_line_dict):
            vcf_line_dict[item] = pos_parts[count]

        sample_format = vcf_line_dict["FORMAT"].split(":")

        for count, item in enumerate(vcf_line_dict):

            if count >= 9:
                genotype = vcf_line_dict[item]

                if  "./." in genotype or genotype == ".":      # "'./.'' for dip, '.' for haploid
                    vcf_line_dict[item] = None

                else:
                    genotype = dict(zip(sample_format, genotype.split(":")))

                    # CONVERT STRINGS TO APPOPRIATE TYPES (INTS, FLOATS, ETC.)
                    if genotype.has_key("GQ"):
                        genotype['GQ'] = float(genotype['GQ'])

                    if genotype.has_key("DP"):
                        genotype['DP'] = int(genotype['DP'])

                    if genotype.has_key("AD"):
                        genotype['AD'] = tuple(int(ad) for ad in genotype['AD'].split(","))

                    if genotype.has_key("PL"):
                        genotype['PL'] = tuple(int(ad) for ad in genotype['PL'].split(","))

                    vcf_line_dict[item] = genotype

        vcf_line_dict['POS'] = int(vcf_line_dict['POS'])

        try:
            vcf_line_dict['QUAL'] = float(vcf_line_dict['QUAL'])
        except ValueError:
            pass

        vcf_line_dict['INFO'] = self.parse_info_field(vcf_line_dict['INFO'])

        return vcf_line_dict.copy()

    def lines_2_dicts(self, chunk):

        vcf_line_dict = self.empty_vcf_line.copy()
        return [self.parse_vcf_line(line, vcf_line_dict) for line in chunk]

    def get_chromosome_lengths(self, regions_to_skip=[]):

        try:
            tbx = pysam.Tabixfile(self.input)  # TODO: create try statement to test that file is actually a VCF
        except:
            print 'Input not Tabix Indexed.'
            sys.exit()

        # PARSE LENGTH INFO FROM HEADER

        chrm_lengths = []
        chrm_lengths_dict = {}
        for line in tbx.header:

            if line.startswith("##contig="):

                chrm_name = re.findall(r'ID=.*,', line)
                chrm_name = chrm_name[0].strip('ID=').strip(',')

                chrm_length = re.findall(r'length=.*>', line)
                chrm_length = int(chrm_length[0].strip('length=').strip('>'))

                if chrm_name in regions_to_skip:
                    print 'skipping', chrm_name
                    continue

                chrm_lengths.append((chrm_name, 1, chrm_length))
                chrm_lengths_dict[chrm_name] = chrm_length

        chrm_lengths = tuple(chrm_lengths)
        tbx.close()

        return chrm_lengths_dict

    def get_slice_indicies(self):
        """Get slice information from VCF file that is tabix indexed (bgzipped). """

        # GENERATE SLICES:
        # Does not yield final partial slice. Not a bug!

        #print [c for c in self.chrms_2_sizes.iteritems()]

        if self.region == [None]:

            # ITERATE OVER CHROMOSOMES (USE ORDERED DICT TO KEEP IN VCF ORDER)
            for chrm, length in self.chrms_2_sizes.iteritems():

                cStart = 0
                cStop = 0
                iCount = 0

                while cStop < length:

                    if iCount == 0:
                        cStart = 1
                        cStop = self.window_size
                        iCount += 1

                    yield (chrm, cStart, cStop)
                    cStart += self.step
                    cStop += self.step

        else:

            chrm, start, stop = self.region

            cStart = 0
            cStop = 0
            iCount = 0

            if self.window_size == None:
                self.window_size = stop - start

            if self.step == None:
                self.step = 0

            while cStop < stop:

                if iCount == 0:
                    cStart = start
                    cStop = start + self.window_size - 1
                    iCount += 1

                yield (chrm, cStart, cStop)

                if self.step == 0:
                    cStart += self.window_size
                else:
                    cStart += self.step

                cStop = cStart + self.window_size + self.step - 1

    def snp_chunks_interator(self, num_snv_per_chunk=3, fillvalue=None):

        def grouper(iterable, n):
            """Yields chunks of SNV of n number of SNVs.
               Properly handles chromsome so that
               chunks don't overlap chromosomes."""

            chrom = None
            chunk = []

            for count, i in enumerate(iterable):

                #s = i.split("\t")[:3]
                current_chrom = i["CHROM"]

                if count == 0:
                    chrom = current_chrom

                if current_chrom != chrom:

                    yield chunk
                    chunk = [i]
                    chrom = current_chrom
                    continue

                if current_chrom == chrom:

                    chunk.append(i)

                    if len(chunk) >= n:
                        yield chunk
                        chunk = []

                chrom = current_chrom

        chunk_iterator = grouper(self.vcf_file_iterator(), n=num_snv_per_chunk)

        for c, i in enumerate(chunk_iterator):

            if len(i) == 0: # handles edge case
                continue

            if len(i) != num_snv_per_chunk:
                padding =  num_snv_per_chunk - len(i)
                i.extend([fillvalue]*padding)

            yield i

    def calc_allele_counts(self, vcf_line_dict, sample_ids=None):

        #allele_counts = defaultdict({0:0.0,1:0.0,2:0.0,3:0.0,4:0.0})
        allele_counts = dict((key, {0: 0.0, 1: 0.0, 2: 0.0, 3: 0.0}) for key in self.populations.keys())

        for population in self.populations.keys():

            for sample_id in self.populations[population]:

                if vcf_line_dict[sample_id] != None:

                    genotype = vcf_line_dict[sample_id]
                    genotype = genotype["GT"].split("/")   # TODO add phased logic

                    if genotype == [".", "."]:
                        continue

                    genotype = [int(item) for item in genotype]

                    for allele in genotype:
                        allele_counts[population][allele] += 1.0

        return allele_counts

<commit_msg>Minor changes<commit_after>#!/usr/bin/env python
# encoding: utf-8

import re
import sys
import gzip
import pysam
#import itertools
import mimetypes
from collections import OrderedDict  #,defaultdict

class VCF(object):
    """docstring for VCF"""

    def __init__(self,
                 input,
                 output=None,
                 populations=None,
                 region=None,
                 window_size=1,
                 step=0,
                 snvs=None):

        super(VCF, self).__init__()

        self.input = input
        self.output = output
        self.step = step
        self.populations = populations
        self.region = region
        self.window_size = window_size
        self.header = self.__extract_header__()
        self.chrms_2_sizes = self.__get_chrm_ids_and_sizes__()
        self.empty_vcf_line = self.make_empty_vcf_ordered_dict()

    def __open_vcf__(self):
        """Open vcf file as gzip or as text."""

        if type(self.input) is file:
            fin = self.input

        elif mimetypes.guess_type(self.input)[-1] is 'gzip':
            fin = gzip.open(self.input, 'rb')

        else:
            fin = open(self.input, 'r')

        return fin

    def __extract_header__(self):
        #print [line.strip() for line in self.__open_vcf__()]

        header = []
        for line in self.__open_vcf__():

            if line.startswith("#") == True:
                header.append(line.strip())
            else:
                break

        return header

    def __get_chrm_ids_and_sizes__(self):
        """ Extract chromosome ids and sizes from vcf file.
            Return as dictionary"""

        chrms_sizes_dict = OrderedDict()
        # with self.__open_vcf__() as fin:

        for line in self.header:

            if line.startswith("##contig"):
                chrm_name = re.findall(r'ID=.*,', line)
                chrm_name = chrm_name[0].strip('ID=').strip(',')

                chrm_length = re.findall(r'length=.*>', line)
                chrm_length = int(chrm_length[0].strip('length=').strip('>'))

                chrms_sizes_dict[chrm_name] = chrm_length
                #break

        return chrms_sizes_dict

    def make_empty_vcf_ordered_dict(self):
        """Open VCF file and read in #CHROM line as an Ordered Dict"""

        header_dict = None
        for line in self.header:
            if line.startswith("#CHROM"):
                header = line.strip("#").strip().split()
                header_dict = OrderedDict([(item, None) for item in header])
                break

        return header_dict

    def process_snp(self, snp_call):
        if snp_call == "0/0":
            return (0,0)

        elif snp_call == "1/1":
            return (1,1)

        elif snp_call == '1/0' or \
               snp_call == '0/1':
            return (0,1)

        # skip multiallelic sites
        else:
            return None

    def process_snp_call(self, snp_call, ref, alt, IUPAC_ambiguities=False):
        """Process VCF genotype fields."""

        # IUPAC ambiguity codes
        IUPAC_dict = {('A', 'C'): 'M',
                      ('A', 'G'): 'R',
                      ('A', 'T'): 'W',
                      ('C', 'G'): 'S',
                      ('C', 'T'): 'Y',
                      ('G', 'T'): 'K',
                      ('A', 'C', 'G'): 'V',
                      ('A', 'C', 'T'): 'H',
                      ('A', 'G', 'T'): 'D',
                      ('C', 'G', 'T'): 'B'}

        if snp_call == None:
            called_base = "-"

        if snp_call["GT"] == "./.":
            called_base = "-"

        else:
            allele1, allele2 = snp_call["GT"].split("/")

            # process "0/0"
            if allele1 == '0' and allele2 == '0':
                called_base = ref

            if allele1 == '1' and allele2 == '1':
                called_base = alt

            # process "0/N"
            if allele1 == '0' and allele2 != '0':

                if IUPAC_ambiguities == False:
                    called_base = 'N'

                else:
                    call = [ref] + [alt.split(',')[int(allele2) - 1]]
                    call.sort()
                    call = tuple(call)
                    called_base = IUPAC_dict[call]

            # process "2/2, 1/2, etc."
            if int(allele1) >= 1 and int(allele2) > 1:

                # deal with homozygotes
                if allele1 == allele2:
                    called_base = alt.split(',')[int(allele1) - 1]

                # deal with heterozygotes
                else:

                    if IUPAC_ambiguities == False:
                        called_base = 'N'

                    else:
                        ref = alt.split(',')[int(allele1) - 1]
                        alt = alt.split(',')[int(allele2) - 1]
                        call = [ref, alt]
                        call.sort()
                        call = tuple(call)
                        called_base = IUPAC_dict[call]

        return called_base

    def count_alleles(self, chunk):

        results = []
        for line in chunk:
            pop_counts = {}

            for pop in self.populations.keys():
                allele_counts = {'REF':0, 'ALT':0}

                for sample in self.populations[pop]:

                    if line[sample] != None:
                        ref, alt = self.process_snp(line[sample]['GT'])
                        allele_counts['REF'] += ref
                        allele_counts['ALT'] += alt

                pop_counts[pop] = allele_counts.copy()

            results.append(pop_counts.copy())

        return results

    def vcf_slice_iterator(self, vcf_bgzipped_file, region):

        tbx = pysam.Tabixfile(vcf_bgzipped_file)

        try:
            vcf_slice = tbx.fetch(*region)
        except ValueError:
            print 'bad vcf slice:', region
            sys.exit()

        else:
            for row in vcf_slice:
                yield self.parse_vcf_line(row, self.empty_vcf_line)
            #return ((chrm, start, stop), chunk)
            #return tuple(row for row in vcf_slice)

    def vcf_file_iterator(self, as_dict=True):

        for line in self.__open_vcf__():
            if line.startswith("#") is not True:
                if as_dict == True:
                    yield self.parse_vcf_line(line, self.empty_vcf_line.copy())
                else:
                    yield line
            else:
                continue

    def parse_info_field(self, info_field):

        info_dict = {}
        for item in info_field.split(';'):
            pair = item.split("=")
            if len(pair) == 2:
                info_dict[pair[0]] = pair[1]   # this could be improved on

        return info_dict

    def get_population_sizes(self, vcfline, populations):

        sample_counts = {}

        for pop in populations.keys():
            sample_count = 0

            for sample in populations[pop]:
                if vcfline[sample] is not None:
                    sample_count += 1

            sample_counts[pop] = sample_count

        return sample_counts

    def parse_vcf_line(self, pos, vcf_line_dict):
        """Read in VCF line and convert it to an OrderedDict"""

        pos_parts = pos.strip().split()
        for count, item in enumerate(vcf_line_dict):
            vcf_line_dict[item] = pos_parts[count]

        sample_format = vcf_line_dict["FORMAT"].split(":")

        for count, item in enumerate(vcf_line_dict):

            if count >= 9:
                genotype = vcf_line_dict[item]

                if  "./." in genotype or genotype == ".":      # "'./.'' for dip, '.' for haploid
                    vcf_line_dict[item] = None

                else:
                    genotype = dict(zip(sample_format, genotype.split(":")))

                    # CONVERT STRINGS TO APPOPRIATE TYPES (INTS, FLOATS, ETC.)
                    if genotype.has_key("GQ"):
                        genotype['GQ'] = float(genotype['GQ'])

                    if genotype.has_key("DP"):
                        genotype['DP'] = int(genotype['DP'])

                    if genotype.has_key("AD"):
                        genotype['AD'] = tuple(int(ad) for ad in genotype['AD'].split(","))

                    if genotype.has_key("PL"):
                        genotype['PL'] = tuple(int(ad) for ad in genotype['PL'].split(","))

                    vcf_line_dict[item] = genotype

        vcf_line_dict['POS'] = int(vcf_line_dict['POS'])

        try:
            vcf_line_dict['QUAL'] = float(vcf_line_dict['QUAL'])
        except ValueError:
            pass

        vcf_line_dict['INFO'] = self.parse_info_field(vcf_line_dict['INFO'])

        return vcf_line_dict.copy()

    def lines_2_dicts(self, chunk):

        vcf_line_dict = self.empty_vcf_line.copy()
        return [self.parse_vcf_line(line, vcf_line_dict) for line in chunk]

    def get_chromosome_lengths(self, regions_to_skip=[]):

        try:
            tbx = pysam.Tabixfile(self.input)  # TODO: create try statement to test that file is actually a VCF
        except:
            print 'Input not Tabix Indexed.'
            sys.exit()

        # PARSE LENGTH INFO FROM HEADER

        chrm_lengths = []
        chrm_lengths_dict = {}
        for line in tbx.header:

            if line.startswith("##contig="):

                chrm_name = re.findall(r'ID=.*,', line)
                chrm_name = chrm_name[0].strip('ID=').strip(',')

                chrm_length = re.findall(r'length=.*>', line)
                chrm_length = int(chrm_length[0].strip('length=').strip('>'))

                if chrm_name in regions_to_skip:
                    print 'skipping', chrm_name
                    continue

                chrm_lengths.append((chrm_name, 1, chrm_length))
                chrm_lengths_dict[chrm_name] = chrm_length

        chrm_lengths = tuple(chrm_lengths)
        tbx.close()

        return chrm_lengths_dict

    def get_slice_indicies(self):
        """Get slice information from VCF file that is tabix indexed (bgzipped). """

        # GENERATE SLICES:
        # Does not yield final partial slice. Not a bug!

        #print [c for c in self.chrms_2_sizes.iteritems()]

        if self.region == [None]:

            # ITERATE OVER CHROMOSOMES (USE ORDERED DICT TO KEEP IN VCF ORDER)
            for chrm, length in self.chrms_2_sizes.iteritems():

                cStart = 0
                cStop = 0
                iCount = 0

                while cStop < length:

                    if iCount == 0:
                        cStart = 1
                        cStop = self.window_size
                        iCount += 1

                    yield (chrm, cStart, cStop)
                    cStart += self.step
                    cStop += self.step

        else:

            chrm, start, stop = self.region

            cStart = 0
            cStop = 0
            iCount = 0

            if self.window_size == None:
                self.window_size = stop - start

            if self.step == None:
                self.step = 0

            while cStop < stop:

                if iCount == 0:
                    cStart = start
                    cStop = start + self.window_size - 1
                    iCount += 1

                yield (chrm, cStart, cStop)

                if self.step == 0:
                    cStart += self.window_size
                else:
                    cStart += self.step

                cStop = cStart + self.window_size + self.step - 1

    def snp_chunks_interator(self, num_snv_per_chunk=3, fillvalue=None):

        def grouper(iterable, n):
            """Yields chunks of SNV of n number of SNVs.
               Properly handles chromsome so that
               chunks don't overlap chromosomes."""

            chrom = None
            chunk = []

            for count, i in enumerate(iterable):

                #s = i.split("\t")[:3]
                current_chrom = i["CHROM"]

                if count == 0:
                    chrom = current_chrom

                if current_chrom != chrom:

                    yield chunk
                    chunk = [i]
                    chrom = current_chrom
                    continue

                if current_chrom == chrom:

                    chunk.append(i)

                    if len(chunk) >= n:
                        yield chunk
                        chunk = []

                chrom = current_chrom

        chunk_iterator = grouper(self.vcf_file_iterator(), n=num_snv_per_chunk)

        for c, i in enumerate(chunk_iterator):

            if len(i) == 0: # handles edge case
                continue

            if len(i) != num_snv_per_chunk:
                padding =  num_snv_per_chunk - len(i)
                i.extend([fillvalue]*padding)

            yield i

    def calc_allele_counts(self, vcf_line_dict, sample_ids=None):

        #allele_counts = defaultdict({0:0.0,1:0.0,2:0.0,3:0.0,4:0.0})
        allele_counts = dict((key, {0: 0.0, 1: 0.0, 2: 0.0, 3: 0.0}) for key in self.populations.keys())

        for population in self.populations.keys():

            for sample_id in self.populations[population]:

                if vcf_line_dict[sample_id] != None:

                    genotype = vcf_line_dict[sample_id]
                    genotype = genotype["GT"].split("/")   # TODO add phased logic

                    if genotype == [".", "."]:
                        continue

                    genotype = [int(item) for item in genotype]

                    for allele in genotype:
                        allele_counts[population][allele] += 1.0

        return allele_counts

<|endoftext|>


<commit_before># Kimi language interpreter in Python 3
# Anjana Vakil
# http://www.github.com/vakila/kimi

import sys
from environment import Environment, standard_env
from errors import *

def tokenize(string):
    '''Take a program as a string, return the tokenized program as a list of strings.

    >>> tokenize("-1")
    [('literal', -1)]

    >>> tokenize("(+ 1 2)")
    [('opening', None), ('symbol', '+'), ('literal', 1), ('literal', 2), ('closing', None)]
    '''
    assert_or_complain(string.count('(') == string.count(')'), "SYNTAX ERROR! Mismatching parentheses.")
    assert_or_throw('(((' not in string, "syntax", 'Incorrect parenthesis use: "(((". Opening parenthesis must be immediately followed by a function.')
    special = ['(',')','"']
    whitespaces = [' ','\n','\t']
    tokens = []
    remaining = string
    while remaining:
        this_char = remaining[0]
        if this_char in whitespaces:
            remaining = remaining[1:]
            continue
        if this_char in ["(", ")"]:
            # the token is this character
            if this_char == "(":
                token_type = 'opening'
                try:
                    next_char = remaining[1]
                except IndexError:
                    throw_error("syntax", 'Incorrect parenthesis use: "(" at end of program.')
                else:
                    if next_char in [")", '"'] or next_char in whitespaces :
                        throw_error("syntax", "Incorrect parenthesis use: " + '"' + this_char + next_char + '". Opening parenthesis must be immediately followed by a function.')
            if this_char == ")":
                token_type = 'closing'
            token_value = None
            remaining = remaining[1:]
        elif this_char == '"':
            # the token is everything until the next "
            endquote_index = remaining[1:].find('"')
            if endquote_index == -1:
                complain_and_die("SYNTAX ERROR! Improper string syntax.")
            endquote_index += 1
            token_value = remaining[1:endquote_index]
            token_type = 'literal'
            remaining = remaining[endquote_index+1:]
        else:
            # the token is everything until the next whitespace or special character
            token_value = ""
            while this_char not in special and this_char not in whitespaces:
                token_value += this_char
                remaining = remaining[1:]
                if not remaining:
                    break
                this_char = remaining[0]
            try:
                # anything that can be converted to int is a literal number
                token_value = int(token_value)
                token_type = "literal"
            except ValueError:
                # everything else is a symbol
                token_type = "symbol"
        tokens.append((token_type, token_value))
    return tokens

def parse(tokens):
    '''Take a list of tokens representing a program, return a tree representing the program's syntax.
    The tree is a nested dictionary, where each dictionary in the tree is an expression.

    Expressions as dictionaries:
    - All contain the key 'type'
    - Based on the value of 'type', the dictionary will also have other keys:
        - type 'apply' (function application) has keys 'operator' (value: expression) and 'arguments' (value: tuple of expressions)
        - type 'symbol' (variable or operator) has key 'name' (value: string representing the variable/operator)
        - type 'literal' (number, string, boolean, ...) has key 'value' (value: the literal)

    >>> parse(tokenize("(+ 1 2)"))
    {'type': 'apply',
     'operator': {'type': 'symbol', 'value': '+'},
     'arguments': ({'type': 'literal', 'value': 1},
                   {'type': 'literal', 'value': 2})}

    '''
    # print("tokens:", tokens)
    if len(tokens) == 0:
        complain_and_die("SYNTAX ERROR! Nothing left to parse.")
    (token_type, token_value) = tokens.pop(0)
    if token_type == 'closing':
        complain_and_die("SYNTAX ERROR! Unexpected ')'.")
    elif token_type == 'opening':
        # print("OPENING")
        operator = parse(tokens)
        arguments = []
        while True:
            if not tokens:
                complain_and_die("SYNTAX ERROR! Unexpected end of program.")
            next_token = tokens[0]
            if next_token[0] == 'closing':
                # print("CLOSING")
                tokens.pop(0)
                # print("tokens:", tokens)
                break
            arguments.append(parse(tokens))
        arguments = tuple(arguments)
        return {'type': 'apply', 'operator': operator, 'arguments': arguments}
    else:
        return {'type': token_type, 'value': token_value}

def special_forms():
    specials = dict()

    def do(args, env):
        do_env = Environment(name="do", outer=env)
        if len(args) == 0:
            throw_error("syntax", "Incorrect use of (do ...): must take at least one argument.")
        result = None
        for a in args:
            result = evaluate(a, do_env)
        return result
    specials['do'] = do

    def lamb(args, env):
        print("\n")
        print("args (" + str(len(args)) + "):", args)
        print("env: ", env.name)
        if len(args) < 2:
            throw_error("syntax", "Incorrect use of (lambda ...): must take at least two arguments (at least one variable and a body).")
        largs = args[:-1]
        lbody = args[-1]
        print("largs (" + str(len(largs)) + "):", largs)
        for l in largs:
            assert_or_throw(l['type'] == 'symbol', "syntax", "Incorrect use of (lambda ...): the anonymous function's variables must be symbols.")
        largs = tuple(la['value'] for la in largs)
        print("lbody:", lbody)
        def anonymous(*arguments):
            print("inside anonymous function")
            print("arguments(" + str(len(arguments)) + "):", arguments)
            if len(arguments) != len(largs):
                throw_error("syntax", "This function takes " + str(len(largs)) + " arguments (" + str(len(arguments)) + " provided).")
            lenv = Environment(name="anon_fn", outer=env, variables=largs, values=arguments)
            return evaluate(lbody, lenv)
        return anonymous
    specials['lambda'] = lamb

    def define(args, env):
        if len(args) != 2:
            throw_error("syntax", "Incorrect use of (define ...): must take exactly two arguments.")
        assert_or_throw(args[0]['type'] == 'symbol', "syntax", "Incorrect use of (define ...): the variable must be a symbol.")
        variable = args[0]['value']
        env.set(variable, evaluate(args[1], env))
    specials['define'] = define

    # def cond(args, env):
    #     pass
    # specials['if'] = cond

    return specials

SPECIALS = special_forms()

def evaluate(expression, environment):
    '''Take an expression and environment as dictionaries.
    Evaluate the expression in the context of the environment, and return the result.

    >>> evaluate(parse(tokenize("(+ 1 2)")), standard_env())
    3
    '''
    print("EVALUATING:", expression)
    expr_type = expression['type']
    print("EXPR_TYPE:", expr_type)
    if expr_type == 'literal':
        return expression['value']
    elif expr_type == 'symbol':
        symbol = expression['value']
        return environment.get(symbol)
    elif expr_type == 'apply':
        operator = expression['operator']
        if operator['type'] == 'symbol' and operator['value'] in SPECIALS:
            return SPECIALS[operator['value']](expression['arguments'], environment)
        fn = evaluate(operator, environment)
        assert_or_throw(callable(fn), "type", 'Trying to call a non-function. Did you use parentheses correctly?')
        return fn(*[evaluate(arg, environment) for arg in expression['arguments']])
    else:
        complain_and_die("PARSING ERROR! Unexpected expression type: " + str(expression) + ".")

def execute(program):
    '''Take a Kimi program as a string. Tokenize the program, parse the tokens into a tree,
    then evaluate the tree. Return the result, or an error message.'''
    return evaluate(parse(tokenize(program)), standard_env())

if __name__ == "__main__":
    try:
        program = sys.argv[1]
    except:
        print("Usage:")
        print("$ python3 kimi.py my_program.kimi")
        print('$ python3 kimi.py "(+ 1 2)"')
    else:
        if program.endswith('.kimi'):
            with open(program, 'r') as f:
                program = f.read()
            print("Evaluating program:")
            print(program)
        print(execute(program))
<commit_msg>Implement if<commit_after># Kimi language interpreter in Python 3
# Anjana Vakil
# http://www.github.com/vakila/kimi

import sys
from environment import Environment, standard_env
from errors import *

def tokenize(string):
    '''Take a program as a string, return the tokenized program as a list of strings.

    >>> tokenize("-1")
    [('literal', -1)]

    >>> tokenize("(+ 1 2)")
    [('opening', None), ('symbol', '+'), ('literal', 1), ('literal', 2), ('closing', None)]
    '''
    assert_or_complain(string.count('(') == string.count(')'), "SYNTAX ERROR! Mismatching parentheses.")
    assert_or_throw('(((' not in string, "syntax", 'Incorrect parenthesis use: "(((". Opening parenthesis must be immediately followed by a function.')
    special = ['(',')','"']
    whitespaces = [' ','\n','\t']
    tokens = []
    remaining = string
    while remaining:
        this_char = remaining[0]
        if this_char in whitespaces:
            remaining = remaining[1:]
            continue
        if this_char in ["(", ")"]:
            # the token is this character
            if this_char == "(":
                token_type = 'opening'
                try:
                    next_char = remaining[1]
                except IndexError:
                    throw_error("syntax", 'Incorrect parenthesis use: "(" at end of program.')
                else:
                    if next_char in [")", '"'] or next_char in whitespaces :
                        throw_error("syntax", "Incorrect parenthesis use: " + '"' + this_char + next_char + '". Opening parenthesis must be immediately followed by a function.')
            if this_char == ")":
                token_type = 'closing'
            token_value = None
            remaining = remaining[1:]
        elif this_char == '"':
            # the token is everything until the next "
            endquote_index = remaining[1:].find('"')
            if endquote_index == -1:
                complain_and_die("SYNTAX ERROR! Improper string syntax.")
            endquote_index += 1
            token_value = remaining[1:endquote_index]
            token_type = 'literal'
            remaining = remaining[endquote_index+1:]
        else:
            # the token is everything until the next whitespace or special character
            token_value = ""
            while this_char not in special and this_char not in whitespaces:
                token_value += this_char
                remaining = remaining[1:]
                if not remaining:
                    break
                this_char = remaining[0]
            try:
                # anything that can be converted to int is a literal number
                token_value = int(token_value)
                token_type = "literal"
            except ValueError:
                # everything else is a symbol
                token_type = "symbol"
        tokens.append((token_type, token_value))
    return tokens

def parse(tokens):
    '''Take a list of tokens representing a program, return a tree representing the program's syntax.
    The tree is a nested dictionary, where each dictionary in the tree is an expression.

    Expressions as dictionaries:
    - All contain the key 'type'
    - Based on the value of 'type', the dictionary will also have other keys:
        - type 'apply' (function application) has keys 'operator' (value: expression) and 'arguments' (value: tuple of expressions)
        - type 'symbol' (variable or operator) has key 'name' (value: string representing the variable/operator)
        - type 'literal' (number, string, boolean, ...) has key 'value' (value: the literal)

    >>> parse(tokenize("(+ 1 2)"))
    {'type': 'apply',
     'operator': {'type': 'symbol', 'value': '+'},
     'arguments': ({'type': 'literal', 'value': 1},
                   {'type': 'literal', 'value': 2})}

    '''
    # print("tokens:", tokens)
    if len(tokens) == 0:
        complain_and_die("SYNTAX ERROR! Nothing left to parse.")
    (token_type, token_value) = tokens.pop(0)
    if token_type == 'closing':
        complain_and_die("SYNTAX ERROR! Unexpected ')'.")
    elif token_type == 'opening':
        # print("OPENING")
        operator = parse(tokens)
        arguments = []
        while True:
            if not tokens:
                complain_and_die("SYNTAX ERROR! Unexpected end of program.")
            next_token = tokens[0]
            if next_token[0] == 'closing':
                # print("CLOSING")
                tokens.pop(0)
                # print("tokens:", tokens)
                break
            arguments.append(parse(tokens))
        arguments = tuple(arguments)
        return {'type': 'apply', 'operator': operator, 'arguments': arguments}
    else:
        return {'type': token_type, 'value': token_value}

def special_forms():
    specials = dict()

    def do(args, env):
        do_env = Environment(name="do", outer=env)
        if len(args) == 0:
            throw_error("syntax", "Incorrect use of (do ...): must take at least one argument.")
        result = None
        for a in args:
            result = evaluate(a, do_env)
        return result
    specials['do'] = do

    def lamb(args, env):
        print("\n")
        print("args (" + str(len(args)) + "):", args)
        print("env: ", env.name)
        if len(args) < 2:
            throw_error("syntax", "Incorrect use of (lambda ...): must take at least two arguments (at least one variable and a body).")
        largs = args[:-1]
        lbody = args[-1]
        print("largs (" + str(len(largs)) + "):", largs)
        for l in largs:
            assert_or_throw(l['type'] == 'symbol', "syntax", "Incorrect use of (lambda ...): the anonymous function's variables must be symbols.")
        largs = tuple(la['value'] for la in largs)
        print("lbody:", lbody)
        def anonymous(*arguments):
            print("inside anonymous function")
            print("arguments(" + str(len(arguments)) + "):", arguments)
            if len(arguments) != len(largs):
                throw_error("syntax", "This function takes " + str(len(largs)) + " arguments (" + str(len(arguments)) + " provided).")
            lenv = Environment(name="anon_fn", outer=env, variables=largs, values=arguments)
            return evaluate(lbody, lenv)
        return anonymous
    specials['lambda'] = lamb

    def define(args, env):
        if len(args) != 2:
            throw_error("syntax", "Incorrect use of (define ...): must take exactly two arguments.")
        assert_or_throw(args[0]['type'] == 'symbol', "syntax", "Incorrect use of (define ...): the variable must be a symbol.")
        variable = args[0]['value']
        env.set(variable, evaluate(args[1], env))
    specials['define'] = define

    def cond(args, env):
        if len(args) != 3:
            throw_error("syntax", "Incorrect use of (if ...): must take exactly three arguments (a test, a success case and a fail case).")
        test = evaluate(args[0], env)
        if test:
            return evaluate(args[1], env)
        else:
            return evaluate(args[2], env)
    specials['if'] = cond

    return specials

SPECIALS = special_forms()

def evaluate(expression, environment):
    '''Take an expression and environment as dictionaries.
    Evaluate the expression in the context of the environment, and return the result.

    >>> evaluate(parse(tokenize("(+ 1 2)")), standard_env())
    3
    '''
    print("EVALUATING:", expression)
    expr_type = expression['type']
    print("EXPR_TYPE:", expr_type)
    if expr_type == 'literal':
        return expression['value']
    elif expr_type == 'symbol':
        symbol = expression['value']
        return environment.get(symbol)
    elif expr_type == 'apply':
        operator = expression['operator']
        if operator['type'] == 'symbol' and operator['value'] in SPECIALS:
            return SPECIALS[operator['value']](expression['arguments'], environment)
        fn = evaluate(operator, environment)
        assert_or_throw(callable(fn), "type", 'Trying to call a non-function. Did you use parentheses correctly?')
        return fn(*[evaluate(arg, environment) for arg in expression['arguments']])
    else:
        complain_and_die("PARSING ERROR! Unexpected expression type: " + str(expression) + ".")

def execute(program):
    '''Take a Kimi program as a string. Tokenize the program, parse the tokens into a tree,
    then evaluate the tree. Return the result, or an error message.'''
    return evaluate(parse(tokenize(program)), standard_env())

if __name__ == "__main__":
    try:
        program = sys.argv[1]
    except:
        print("Usage:")
        print("$ python3 kimi.py my_program.kimi")
        print('$ python3 kimi.py "(+ 1 2)"')
    else:
        if program.endswith('.kimi'):
            with open(program, 'r') as f:
                program = f.read()
            print("Evaluating program:")
            print(program)
        print(execute(program))
<|endoftext|>


<commit_before><commit_msg>Delete afm.py<commit_after><|endoftext|>


<commit_before>#!/usr/bin/env python3
# Artemis Mission Compiler: https://github.com/hunternet93/AMC
# by Isaac Smith
# Released under the MIT license, please see LICENSE for more info

import sys, random, copy
import xml.etree.cElementTree as ET
from bunch import *

class AddList(list):
    def add(self, *items): [self.append(item) for item in items]

Mission = lambda: Bunch(tag = 'mission', start = AddList(), events = AddList())
Event = lambda name, f='all': Bunch(name = name, tag = 'event', fire = f, conditions = AddList(), actions = AddList())
Tag = lambda tag, **args: Bunch(tag = tag, **args)

SetVariable = lambda name, value: Bunch(tag = 'set_variable', name = name, value = value)
SetTimer = lambda name, seconds: Bunch(tag = 'set_timer', name = name, seconds = seconds)
SetObjectProperty = lambda name, property, value: Bunch(tag = 'set_object_property', name = name, property = property, value=value)
ClearAI = lambda name: Bunch(tag = 'clear_ai', name = name)
AddAI = lambda name, type, **args: Bunch(tag = 'add_ai', name = name, type = type, **args)
EndMission = lambda: Bunch(tag = 'end_mission')
BigMessage = lambda title, subtitle1 = '', subtitle2 = '': Bunch(tag = 'big_message', title = title, subtitle1 = subtitle1, subtitle2 = subtitle2)
IncomingCommsText = lambda from_, body: Bunch(tag = 'incoming_comms_text', from_ = from_, body = body)
Destroy = lambda name: Bunch(tag = 'destroy', name = name)

IfVariable = lambda name, comparator, value: Bunch(tag = 'if_variable', name = name, comparator = comparator, value = value)
IfTimerFinished = lambda name: Bunch(tag = 'if_timer_finished', name = name)
IfExists = lambda name: Bunch(tag = 'if_exists', name = name)
IfNotExists = lambda name: Bunch(tag = 'if_not_exists', name = name)
IfInsideBox = lambda name, leastX, leastZ, mostX, mostZ: Bunch(tag = 'if_inside_box', name = name, leastX = leastX, leastZ = leastZ, mostX = mostX, mostZ = mostZ)
IfOutsideBox = lambda name, leastX, leastZ, mostX, mostZ: Bunch(tag = 'if_outside_box', name = name, leastX = leastX, leastZ = leastZ, mostX = mostX, mostZ = mostZ)
IfInsideSphere = lambda name, centerX, centerY, centerZ, radius: Bunch(tag = 'if_inside_sphere', name = name, centerX = centerX, centerY = centerY, centerZ = centerZ, radius = radius)
IfOutsideSphere = lambda name, centerX, centerY, centerZ, radius: Bunch(tag = 'if_outside_sphere', name = name, centerX = centerX, centerY = centerY, centerZ = centerZ, radius = radius)
IfDistance = lambda name1, name2, comparator, value: Bunch(tag = 'if_distance', name1 = name1, name2 = name2, comparator = comparator, value = value)
IfDamconMembers = lambda team_index, comparator, value: Bunch(tag = 'if_damcom_members', team_index = team_index, comparator = comparator, value = value)
IfDifficulty = lambda comparator, value: Bunch(tag = 'if_difficulty', comparator = comparator, value = value)
IfFleetCount = lambda fleetnumber, comparator, value: Bunch(tag = 'if_fleet_count', fleetnumber = fleetnumber, comparator = comparator, value = value)
IfEnemyCount = lambda comparator, value: Bunch(tag = 'if_fleet_count', comparator = comparator, value = value)
IfDocked = lambda name: Bunch(tag = 'if_docked', name = name)
IfPlayerIsTargeting = lambda name: Bunch(tag = 'if_player_is_targeting', name = name)
IfObjectProperty = lambda name, property, comparator, value: Bunch(tag = 'if_object_property', name = name, property = property, comparator = comparator, value = value)

Player = lambda name, **args: Bunch(name = name, tag = 'create', type = 'player', **args)
Enemy = lambda name, **args: Bunch(name = name, tag = 'create', type = 'enemy', **args)
Station = lambda name, **args: Bunch(name = name, tag = 'create', type = 'station', **args)
Neutral = lambda name, **args: Bunch(name = name, tag = 'create', type = 'neutral', **args)
Anomaly = lambda name, **args: Bunch(name = name, tag = 'create', type = 'anomaly', **args)
BlackHole = lambda name, **args: Bunch(name = name, tag = 'create', type = 'blackhole', **args)
Monster = lambda name, **args: Bunch(name = name, tag = 'create', type = 'monster', **args)
GenericMesh = lambda name, **args: Bunch(name = name, tag = 'create', type = 'genericmesh', **args)
Whale = lambda name, **args: Bunch(name = name, tag = 'create', type = 'whale', **args)
Nebulas = lambda **args: Bunch(tag = 'create', type = 'nebulas', **args)
Asteroids = lambda **args: Bunch(tag = 'create', type = 'asteroids', **args)
Mines = lambda **args: Bunch(tag = 'create', type = 'mines', **args)

def create_tag(item, parent):
    tag = ET.SubElement(parent, item.pop('tag'))
    if item.get('body'):
        tag.text = item.pop('body')

    for argument in item.items():
        tag.set(str(argument[0]).replace('_', ""), str(argument[1]))

def create_entity(item, parent):
    entity = ET.SubElement(parent, item.pop('tag'))
    entity.set('type', item.pop('type'))

    if item.get('properties'):
        properties = item.pop('properties')
    else:
        properties = False

    if item.get('ai'):
        ai = item.pop('ai')
    else:
        ai = False

    for attribute in item.items():
        entity.set(str(attribute[0]), str(attribute[1]))

    if ai:
        ET.SubElement(parent, 'clear_ai').set('name', name)
        for block in ai:
            add_ai = ET.SubElement(parent, 'add_ai')
            add_ai.set('name', name)
            for argument in command.items():
                add_ai.set(str(argument[0]), str(argument[1]))

def parse_items(items, parent):
    for item in items:
        if item.tag == 'create':
            create_entity(item, parent)
        else:
            create_tag(item, parent)

def parse_event(event, parent):
    tag = ET.SubElement(parent, 'event')

    if event['fire'] == 'once':
        event.conditions.add(IfVariable(event.name + '_triggered', '!=', 1))
        event.actions.add(SetVariable(event.name + '_triggered', 1))

    for condition in event.conditions:
        create_tag(condition, tag)

    parse_items(event.actions, tag)

def parse(mission):
    root = ET.Element('mission')
    root.set('version', '1')

    start = ET.SubElement(root, 'start')
    parse_items(mission.start, start)

    if mission.get('events'):
        for event in mission['events']:
            parse_event(event, root)

    return ET.tostring(root)

if __name__ == "__main__":
    if not len(sys.argv) >= 2:
        print('Usage: amc.py mission.amc.py <mission.xml>')
        quit()
    else:
        mission_file = open(sys.argv[1], "r")
        if len(sys.argv) == 3:
            output = open(sys.argv[2], "w")
        else:
            output = sys.stdout

    mission = Mission()
    start = mission.start
    events = mission.events
    exec(mission_file.read(), globals())
    output.write(parse(mission).decode())
    output.write("\n")
    output.close()
<commit_msg>Added additional action objects<commit_after>#!/usr/bin/env python3
# Artemis Mission Compiler: https://github.com/hunternet93/AMC
# by Isaac Smith
# Released under the MIT license, please see LICENSE for more info

import sys, random, copy
import xml.etree.cElementTree as ET
from bunch import *

class AddList(list):
    def add(self, *items): [self.append(item) for item in items]

Mission = lambda: Bunch(tag = 'mission', start = AddList(), events = AddList())
Event = lambda name, f='all': Bunch(name = name, tag = 'event', fire = f, conditions = AddList(), actions = AddList())
Tag = lambda tag, **args: Bunch(tag = tag, **args)

SetVariable = lambda name, value: Bunch(tag = 'set_variable', name = name, value = value)
SetTimer = lambda name, seconds: Bunch(tag = 'set_timer', name = name, seconds = seconds)
SetObjectProperty = lambda name, property, value: Bunch(tag = 'set_object_property', name = name, property = property, value = value)
AddToObjectProperty = lambda name, property, value: Bunch(tag = 'addto_object_property', name = name, property = property, value = value)
CopyObjectProperty = lambda name1, name2, value: Bunch(tag = 'copy_object_property', name1 = name1, name2 = name2, value = value)
SetFleetProperty = lambda fleetIndex, property, value: Bunch(tag = 'set_fleet_property', fleetIndex = fleetIndex, property = property, value = value)
ClearAI = lambda name: Bunch(tag = 'clear_ai', name = name)
AddAI = lambda name, type, **args: Bunch(tag = 'add_ai', name = name, type = type, **args)
EndMission = lambda: Bunch(tag = 'end_mission')
BigMessage = lambda title, subtitle1 = '', subtitle2 = '': Bunch(tag = 'big_message', title = title, subtitle1 = subtitle1, subtitle2 = subtitle2)
IncomingCommsText = lambda from_, body: Bunch(tag = 'incoming_comms_text', from_ = from_, body = body)
Destroy = lambda name: Bunch(tag = 'destroy', name = name)
Log = lambda text: bunch(tag = 'log', text = text)

IfVariable = lambda name, comparator, value: Bunch(tag = 'if_variable', name = name, comparator = comparator, value = value)
IfTimerFinished = lambda name: Bunch(tag = 'if_timer_finished', name = name)
IfExists = lambda name: Bunch(tag = 'if_exists', name = name)
IfNotExists = lambda name: Bunch(tag = 'if_not_exists', name = name)
IfInsideBox = lambda name, leastX, leastZ, mostX, mostZ: Bunch(tag = 'if_inside_box', name = name, leastX = leastX, leastZ = leastZ, mostX = mostX, mostZ = mostZ)
IfOutsideBox = lambda name, leastX, leastZ, mostX, mostZ: Bunch(tag = 'if_outside_box', name = name, leastX = leastX, leastZ = leastZ, mostX = mostX, mostZ = mostZ)
IfInsideSphere = lambda name, centerX, centerY, centerZ, radius: Bunch(tag = 'if_inside_sphere', name = name, centerX = centerX, centerY = centerY, centerZ = centerZ, radius = radius)
IfOutsideSphere = lambda name, centerX, centerY, centerZ, radius: Bunch(tag = 'if_outside_sphere', name = name, centerX = centerX, centerY = centerY, centerZ = centerZ, radius = radius)
IfDistance = lambda name1, name2, comparator, value: Bunch(tag = 'if_distance', name1 = name1, name2 = name2, comparator = comparator, value = value)
IfDamconMembers = lambda team_index, comparator, value: Bunch(tag = 'if_damcom_members', team_index = team_index, comparator = comparator, value = value)
IfDifficulty = lambda comparator, value: Bunch(tag = 'if_difficulty', comparator = comparator, value = value)
IfFleetCount = lambda fleetnumber, comparator, value: Bunch(tag = 'if_fleet_count', fleetnumber = fleetnumber, comparator = comparator, value = value)
IfEnemyCount = lambda comparator, value: Bunch(tag = 'if_fleet_count', comparator = comparator, value = value)
IfDocked = lambda name: Bunch(tag = 'if_docked', name = name)
IfPlayerIsTargeting = lambda name: Bunch(tag = 'if_player_is_targeting', name = name)
IfObjectProperty = lambda name, property, comparator, value: Bunch(tag = 'if_object_property', name = name, property = property, comparator = comparator, value = value)

Player = lambda name, **args: Bunch(name = name, tag = 'create', type = 'player', **args)
Enemy = lambda name, **args: Bunch(name = name, tag = 'create', type = 'enemy', **args)
Station = lambda name, **args: Bunch(name = name, tag = 'create', type = 'station', **args)
Neutral = lambda name, **args: Bunch(name = name, tag = 'create', type = 'neutral', **args)
Anomaly = lambda name, **args: Bunch(name = name, tag = 'create', type = 'anomaly', **args)
BlackHole = lambda name, **args: Bunch(name = name, tag = 'create', type = 'blackhole', **args)
Monster = lambda name, **args: Bunch(name = name, tag = 'create', type = 'monster', **args)
GenericMesh = lambda name, **args: Bunch(name = name, tag = 'create', type = 'genericmesh', **args)
Whale = lambda name, **args: Bunch(name = name, tag = 'create', type = 'whale', **args)
Nebulas = lambda **args: Bunch(tag = 'create', type = 'nebulas', **args)
Asteroids = lambda **args: Bunch(tag = 'create', type = 'asteroids', **args)
Mines = lambda **args: Bunch(tag = 'create', type = 'mines', **args)

def create_tag(item, parent):
    tag = ET.SubElement(parent, item.pop('tag'))
    if item.get('body'):
        tag.text = item.pop('body')

    for argument in item.items():
        tag.set(str(argument[0]).replace('_', ""), str(argument[1]))

def create_entity(item, parent):
    entity = ET.SubElement(parent, item.pop('tag'))
    entity.set('type', item.pop('type'))

    if item.get('properties'):
        properties = item.pop('properties')
    else:
        properties = False

    if item.get('ai'):
        ai = item.pop('ai')
    else:
        ai = False

    for attribute in item.items():
        entity.set(str(attribute[0]), str(attribute[1]))

    if ai:
        ET.SubElement(parent, 'clear_ai').set('name', name)
        for block in ai:
            add_ai = ET.SubElement(parent, 'add_ai')
            add_ai.set('name', name)
            for argument in command.items():
                add_ai.set(str(argument[0]), str(argument[1]))

def parse_items(items, parent):
    for item in items:
        if item.tag == 'create':
            create_entity(item, parent)
        else:
            create_tag(item, parent)

def parse_event(event, parent):
    tag = ET.SubElement(parent, 'event')

    if event['fire'] == 'once':
        event.conditions.add(IfVariable(event.name + '_triggered', '!=', 1))
        event.actions.add(SetVariable(event.name + '_triggered', 1))

    for condition in event.conditions:
        create_tag(condition, tag)

    parse_items(event.actions, tag)

def parse(mission):
    root = ET.Element('mission')
    root.set('version', '1')

    start = ET.SubElement(root, 'start')
    parse_items(mission.start, start)

    if mission.get('events'):
        for event in mission['events']:
            parse_event(event, root)

    return ET.tostring(root)

if __name__ == "__main__":
    if not len(sys.argv) >= 2:
        print('Usage: amc.py mission.amc.py <mission.xml>')
        quit()
    else:
        mission_file = open(sys.argv[1], "r")
        if len(sys.argv) == 3:
            output = open(sys.argv[2], "w")
        else:
            output = sys.stdout

    mission = Mission()
    start = mission.start
    events = mission.events
    exec(mission_file.read(), globals())
    output.write(parse(mission).decode())
    output.write("\n")
    output.close()
<|endoftext|>


<commit_before>#!/usr/bin/env python
import sys, os, zipfile
import options
from db import Db

MAX_PACKSIZE = 1024*1024*1024

OPTSPEC = """
kurt.py [options] <path>
--
 Options:
o,output= path to the output directory
v,verbose increase log output (can be used more than once)
"""

def get_files(path):
    try:
        files = os.listdir(path)
    except OSError:
        sys.stderr.write("Couldn't traverse %s\n" % path)
        files = []

    for file in sorted(files):
        if os.path.isfile(os.path.join(path, file)):
            file_path = os.path.join(path, file)
            try:
                file_stats = os.stat(file_path)
                file_size = file_stats.st_size
            except OSError:
                file_size = 0
            yield (file_path, file_size)
        if os.path.isdir(os.path.join(path, file)):
            for file_path, file_size in get_files(os.path.join(path, file)):
                yield (file_path, file_size)

def get_packs(path, known_files):
    singles = []

    packsize = 0
    paths = []
    for file_path, file_size in get_files(path):
        if file_path[len(path)+1:] in known_files:
            continue

        if packsize + file_size <= MAX_PACKSIZE:
            paths.append(file_path)
            packsize += file_size
        elif file_size > MAX_PACKSIZE:
            singles.append(([file_path], file_size))
        else:
            yield (paths, packsize)
            packsize = 0
            paths = []

            for single in singles:
                yield single
            singles = []
    if len(paths) > 0:
        yield (paths, packsize)

def main():
    o = options.Options(OPTSPEC)
    (opt, flags, extra) = o.parse(sys.argv[1:])

    if not extra:
        o.fatal('no path given')

    PATH = extra[0]
    if not os.path.exists(PATH):
        o.fatal('path <%s> does not exist' % PATH)

    if opt.output and not os.path.exists(opt.output):
        o.fata('output-path <%s> does not exist' % opt.output)

    if opt.output:
        metapath = os.path.join(opt.output, 'meta')
    else:
        metapath = 'meta'

    db = Db(os.path.join(metapath, 'kurt.db'))

    known_files, maxpack = db.get_files()

    for index, pack in enumerate(get_packs(PATH, known_files)):
        zipname = "pack%04d.zip" % index
        if opt.output:
            zippath = os.path.join(opt.output, zipname)
        else:
            zippath = zipname
        f = open(zippath, 'wb')
        z = zipfile.ZipFile(f, 'w')

        paths, packsize = pack
        if opt.verbose:
            print "Pack %d (%dB)" % (maxpack+index+1, packsize)
        for path in paths:
            if opt.verbose:
                print "  %s" % path
            pathlen = len(PATH)
            try:
                z.write(path, path[pathlen+1:])
                db.add(path[pathlen+1:], maxpack+index+1)
                file_stats = os.stat(path)
                file_size = file_stats.st_size
                db.add_meta(path[pathlen+1:], "size", file_size)
            except OSError:
                pass
            except IOError:
                pass

        z.close()
        f.close()
    db.commit()

if __name__ == "__main__":
    main()

<commit_msg>Create metapath if not exists<commit_after>#!/usr/bin/env python
import sys, os, zipfile
import options
from db import Db

MAX_PACKSIZE = 1024*1024*1024

OPTSPEC = """
kurt.py [options] <path>
--
 Options:
o,output= path to the output directory
v,verbose increase log output (can be used more than once)
"""

def get_files(path):
    try:
        files = os.listdir(path)
    except OSError:
        sys.stderr.write("Couldn't traverse %s\n" % path)
        files = []

    for file in sorted(files):
        if os.path.isfile(os.path.join(path, file)):
            file_path = os.path.join(path, file)
            try:
                file_stats = os.stat(file_path)
                file_size = file_stats.st_size
            except OSError:
                file_size = 0
            yield (file_path, file_size)
        if os.path.isdir(os.path.join(path, file)):
            for file_path, file_size in get_files(os.path.join(path, file)):
                yield (file_path, file_size)

def get_packs(path, known_files):
    singles = []

    packsize = 0
    paths = []
    for file_path, file_size in get_files(path):
        if file_path[len(path)+1:] in known_files:
            continue

        if packsize + file_size <= MAX_PACKSIZE:
            paths.append(file_path)
            packsize += file_size
        elif file_size > MAX_PACKSIZE:
            singles.append(([file_path], file_size))
        else:
            yield (paths, packsize)
            packsize = 0
            paths = []

            for single in singles:
                yield single
            singles = []
    if len(paths) > 0:
        yield (paths, packsize)

def main():
    o = options.Options(OPTSPEC)
    (opt, flags, extra) = o.parse(sys.argv[1:])

    if not extra:
        o.fatal('no path given')

    PATH = extra[0]
    if not os.path.exists(PATH):
        o.fatal('path <%s> does not exist' % PATH)

    if opt.output and not os.path.exists(opt.output):
        o.fata('output-path <%s> does not exist' % opt.output)

    if opt.output:
        metapath = os.path.join(opt.output, 'meta')
    else:
        metapath = 'meta'
    if not os.path.exists(metapath):
        os.makedirs(metapath)

    db = Db(os.path.join(metapath, 'kurt.db'))

    known_files, maxpack = db.get_files()

    for index, pack in enumerate(get_packs(PATH, known_files)):
        zipname = "pack%04d.zip" % index
        if opt.output:
            zippath = os.path.join(opt.output, zipname)
        else:
            zippath = zipname
        f = open(zippath, 'wb')
        z = zipfile.ZipFile(f, 'w')

        paths, packsize = pack
        if opt.verbose:
            print "Pack %d (%dB)" % (maxpack+index+1, packsize)
        for path in paths:
            if opt.verbose:
                print "  %s" % path
            pathlen = len(PATH)
            try:
                z.write(path, path[pathlen+1:])
                db.add(path[pathlen+1:], maxpack+index+1)
                file_stats = os.stat(path)
                file_size = file_stats.st_size
                db.add_meta(path[pathlen+1:], "size", file_size)
            except OSError:
                pass
            except IOError:
                pass

        z.close()
        f.close()
    db.commit()

if __name__ == "__main__":
    main()

<|endoftext|>