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

"""
vcf2oneliners.py

Created by Nick Crawford on 2012-24-04.
Copyright (c) 2012

The author may be contacted at ngcrawford@gmail.com
"""

import os
import sys
import math
import argparse
import numpy as np
from fstats import fstats
from itertools import combinations
from collections import OrderedDict

def get_args():
    """Parse sys.argv"""
    parser = argparse.ArgumentParser()
    parser.add_argument('-i','--input', required=True,
                        help='Path to VCF file.')
    parser.add_argument('-o','--output',
                        help='Path to output csv file. If path is not set defaults to STDOUT.')
    parser.add_argument('-c','--cores', required=True,
                        help='Number of cores to use.')
    parser.add_argument('-p','--populations', nargs='+',
                        help='Names of populations and samples. The format is: "PopName:sample1,sample2,sample3,etc..."')

    parser.add_argument('-w','--window-size',default=None, type=int,
                        help='Size of the window in which to calculate pairwise F-staticstics')
    parser.add_argument("-f",'--f-statistic',required=True,
                        choices=['Gst_est', 'G_prime_st_est', 'G_double_prime_st_est', 'D_est'])

    args = parser.parse_args()

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

    args.populations = populations_dict

    return args

class VCF(object):
    """docstring for VCF"""
    def __init__(self):
        super(VCF, self).__init__()

        self.header = None
        self.__header_dict__ = None

        self.sample_format = None
        self.sample_format_dict = None

        self.populations = None

    def parse_individual_snps(self, vcf, fout, stat_id):

        vcf = open(vcf,'rU')

        if fout == None:
            fout = sys.stdout
        else:
            fout = open(fout,'w')

        # SETUP NAMED TUPLE TO STORE INFO FROM A SINGLE BASE
        field_labels = []

        # SETUP COUNTERS
        current_base = None
        line_count = None

        # PARSE VCF FIlE
        snp_count = 0

        for line in vcf:

            # SKIP HEADER
            if line.startswith("#CHROM"):
                self.header = line.strip("#").strip().split()
                self.__header_dict__ = OrderedDict([(item,None) for item in self.header])
                line_count = 0

            # START PROCESSING ALIGNED BASES
            if line_count > 0:
                vcf_line_dict = self.parse_vcf_line(line)

                allele_counts = self.calc_allele_counts(vcf_line_dict)
                f_statistics = self.calc_fstats(allele_counts)
                info = self.parse_info(vcf_line_dict["INFO"])
                formated_data, header = self.__write_to_outfiles__(vcf_line_dict["CHROM"],\
                                    vcf_line_dict["POS"],info["DP"], stat_id, f_statistics)

                if line_count == 1:
                    fout.write(header+"\n")
                    fout.write(formated_data+"\n")
                else:
                    fout.write(formated_data+"\n")

            if line_count >= 0:
                line_count += 1

    def parse_info(self,info_field):

        info = []
        for item in info_field.split(','): # TO DO: comma should be ';'
            pair = item.split("=")
            if len(pair) == 2:
                info.append(pair)

        info_dict = dict(info)
        return info_dict

    def parse_vcf_line(self, line):

        line_parts = line.strip().split()
        vcf_line_dict = self.__header_dict__

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

        if self.sample_format == None:
            self.sample_format = vcf_line_dict["FORMAT"].split(":")
            # self.sample_format_dict = dict([(item,None) for item in self.sample_format])

        for count, item in enumerate(vcf_line_dict):

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

                if genotype == "./.":
                    vcf_line_dict[item] = None

                else:
                    genotype = dict(zip(self.sample_format,genotype.split(":")))
                    vcf_line_dict[item] = genotype

        return vcf_line_dict

    def calc_allele_counts(self, vcf_line_dict):

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

        for population in self.populations.keys():

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

            for sample_id in self.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

        return allele_counts

    def calc_fstats(self, 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 = self.populations.fromkeys(self.populations.keys(),None)
        for population in allele_counts.keys():
            counts =  allele_counts[population].values()
            freqs =  counts/np.sum(counts,dtype=float)
            allele_freqs_dict[population] = freqs

        allele_freqs = allele_freqs_dict.values()

        # CACULATE PAIRWISE F-STATISTICS
        pairwise_results = {}
        for population_pair in combinations(self.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 __write_to_outfiles__(self, chrm, pos, depth, stat_id, f_statistics):

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

            if f_statistics[pair][stat_id] == np.nan:
                value = 'NA'
            else:
                value = str(f_statistics[pair][stat_id])
            line += "," + value

        header = 'chrm,pos,total_depth,' + ','.join([left+"-"+right for left, right in f_statistics.keys()]) + "," + stat_id
        return (line, header)

    def slidingWindow(self, vcf, window_size=1000):
        """Generator function that yields non overlapping chunks of VCF lines."""

        # TO DO: add overlapping increment
        chrm_id = None
        start = 0
        stop = window_size
        chunk = []

        for count, line in enumerate(vcf):

            # SKIP HEADER
            if line.startswith("#"):
                continue

            # GET CHRM AND POS
            line_parts = line.strip().split()
            chrm, pos = line_parts[:2]
            pos = int(pos)

            # UPDATE CHRM ID FOR INITAL SITE
            if chrm_id == None:
                chrm_id = chrm

            # REZERO AT NEW CHRM
            #   AND YIELD CHUNK
            if chrm_id != chrm:
                yield chunk
                stop = window_size
                chrm_id = chrm
                chunk = [line]
                continue

            # UPDATE CHUNK
            if pos < stop:
                chunk.append(line)

            # YIELD CHUNK IF CRURRENT
            #   POS EXCEEDS STOP
            if pos >= stop:
                yield chunk
                stop += window_size
                chunk = [line]

    def set_header(self, vcf_path):
        vcf_file = open(vcf_path,'rU')
        for line in vcf_file:
            if line.startswith("#CHROM"):
                self.header = line.strip("#").strip().split()
                self.__header_dict__ = OrderedDict([(item,None) for item in self.header])

        vcf_file.close()

    def process_window(self, window, args):

        Hs_est_list = []
        Ht_est_list = []

        chrm = None
        pos = None
        depth = None

        for count, line in enumerate(window):
            vcf_line_dict = self.parse_vcf_line(line)
            if count == 0:
                chrm = vcf_line_dict["CHROM"]
                pos = vcf_line_dict['POS']
                info = self.parse_info(vcf_line_dict["INFO"])
                depth = info["DP"]

            allele_counts = self.calc_allele_counts(vcf_line_dict)
            f_statistics = self.calc_fstats(allele_counts)

            Hs_est = f_statistics[('MAR', 'CAP')]['Hs_est']
            Hs_est_list.append(Hs_est)

            Ht_est = f_statistics[('MAR', 'CAP')]['Ht_est']
            Ht_est_list.append(Ht_est)

        # Remove uninformative SNPs (nans)
        Hs_est_list = [item for item in Hs_est_list if not math.isnan(item)]
        Ht_est_list = [item for item in Ht_est_list if not math.isnan(item)]

        multilocus_f_statistics = {('MAR', 'CAP'): None}
        if len(Hs_est_list) != 0 or len(Ht_est_list)!= 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[('MAR', 'CAP')] = values_dict

        self.__write_to_outfiles__( chrm, pos, depth, args.f_statistic, multilocus_f_statistics)

if __name__ == '__main__':

    args = get_args()

    vcf = VCF()
    vcf_file = args.input
    vcf.set_header(vcf_file)
    vcf.populations = args.populations

    if args.window_size != None:

        vcf_file = open(vcf_file,'rU')
        for window in vcf.slidingWindow(vcf_file,2000):
            vcf.process_window(window, args)

    else:
        vcf.populations = args.populations
        vcf.parse_individual_snps(args.input, args.output, args.f_statistic)

<commit_msg>Needs further cleaning up.<commit_after>#!/usr/bin/env python
# encoding: utf-8

"""
vcf2oneliners.py

Created by Nick Crawford on 2012-24-04.
Copyright (c) 2012

The author may be contacted at ngcrawford@gmail.com
"""

import os
import sys
import math
import argparse
import copy_reg
import types
import multiprocessing
import numpy as np
from fstats import fstats
from itertools import combinations, izip_longest
from collections import OrderedDict

def get_args():
    """Parse sys.argv"""
    parser = argparse.ArgumentParser()
    parser.add_argument('-i','--input', required=True,
                        help='Path to VCF file.')
    parser.add_argument('-o','--output',
                        help='Path to output csv file. If path is not set defaults to STDOUT.')
    parser.add_argument('-c','--cores', required=True,
                        help='Number of cores to use.')
    parser.add_argument('-p','--populations', nargs='+',
                        help='Names of populations and samples. The format is: "PopName:sample1,sample2,sample3,etc..."')

    parser.add_argument('-w','--window-size',default=None, type=int,
                        help='Size of the window in which to calculate pairwise F-staticstics')
    parser.add_argument("-f",'--f-statistic',required=True,
                        choices=['Gst_est', 'G_prime_st_est', 'G_double_prime_st_est', 'D_est'])

    args = parser.parse_args()

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

    args.populations = populations_dict

    return args

def _pickle_method(method):
    func_name = method.im_func.__name__
    obj = method.im_self
    cls = method.im_class

    return _unpickle_method, (func_name, obj, cls)

def _unpickle_method(func_name, obj, cls):
    for cls in cls.mro():
        try:
            func = cls.__dict__[func_name]
        except KeyError:
            pass
        else:
            break
    return func.__get__(obj, cls)

class VCF(object):
    """docstring for VCF"""
    def __init__(self):
        super(VCF, self).__init__()

        self.header = None
        self.__header_dict__ = None

        self.sample_format = None
        self.sample_format_dict = None

        self.populations = None
        self.f_statistic  = None
        self.vcf_file = None
        self.window_size = None

    def parse_individual_snps(self, vcf, fout, stat_id):

        vcf = open(vcf,'rU')

        if fout == None:
            fout = sys.stdout
        else:
            fout = open(fout,'w')

        # SETUP NAMED TUPLE TO STORE INFO FROM A SINGLE BASE
        field_labels = []

        # SETUP COUNTERS
        current_base = None
        line_count = None

        # PARSE VCF FIlE
        snp_count = 0

        for line in vcf:

            # SKIP HEADER
            if line.startswith("#CHROM"):
                self.header = line.strip("#").strip().split()
                self.__header_dict__ = OrderedDict([(item,None) for item in self.header])
                line_count = 0

            # START PROCESSING ALIGNED BASES
            if line_count > 0:
                vcf_line_dict = self.parse_vcf_line(line)

                allele_counts = self.calc_allele_counts(vcf_line_dict)
                f_statistics = self.calc_fstats(allele_counts)
                info = self.parse_info(vcf_line_dict["INFO"])
                formated_data, header = self.__write_to_outfiles__(vcf_line_dict["CHROM"],\
                                    vcf_line_dict["POS"],info["DP"], stat_id, f_statistics)

                if line_count == 1:
                    fout.write(header+"\n")
                    fout.write(formated_data+"\n")
                else:
                    fout.write(formated_data+"\n")

            if line_count >= 0:
                line_count += 1

    def parse_info(self,info_field):

        info = []
        for item in info_field.split(','): # TO DO: comma should be ';'
            pair = item.split("=")
            if len(pair) == 2:
                info.append(pair)

        info_dict = dict(info)
        return info_dict

    def parse_vcf_line(self, line):

        line_parts = line.strip().split()
        vcf_line_dict = self.__header_dict__

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

        if self.sample_format == None:
            self.sample_format = vcf_line_dict["FORMAT"].split(":")
            # self.sample_format_dict = dict([(item,None) for item in self.sample_format])

        for count, item in enumerate(vcf_line_dict):

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

                if genotype == "./.":
                    vcf_line_dict[item] = None

                else:
                    genotype = dict(zip(self.sample_format,genotype.split(":")))
                    vcf_line_dict[item] = genotype

        return vcf_line_dict

    def calc_allele_counts(self, vcf_line_dict):

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

        for population in self.populations.keys():

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

            for sample_id in self.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

        return allele_counts

    def calc_fstats(self, 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 = self.populations.fromkeys(self.populations.keys(),None)
        for population in allele_counts.keys():
            counts =  allele_counts[population].values()
            freqs =  counts/np.sum(counts,dtype=float)
            allele_freqs_dict[population] = freqs

        allele_freqs = allele_freqs_dict.values()

        # CACULATE PAIRWISE F-STATISTICS
        pairwise_results = {}
        for population_pair in combinations(self.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 __write_to_outfiles__(self, chrm, pos, depth, stat_id, f_statistics):

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

            if f_statistics[pair][stat_id] == np.nan:
                value = 'NA'
            else:
                value = str(f_statistics[pair][stat_id])
            line += "," + value

        header = 'chrm,pos,total_depth,' + ','.join([left+"-"+right for left, right in f_statistics.keys()]) + "," + stat_id
        return (line, header)

    def set_header(self, vcf_path):
        vcf_file = open(vcf_path,'rU')
        for line in vcf_file:
            if line.startswith("#CHROM"):
                self.header = line.strip("#").strip().split()
                self.__header_dict__ = OrderedDict([(item,None) for item in self.header])

        vcf_file.close()

def grouper(n, iterable, padvalue=None):
    """grouper(3, 'abcdefg', 'x') -->
    ('a','b','c'), ('d','e','f'), ('g','x','x')"""

    return izip_longest(*[iter(iterable)]*n, fillvalue=padvalue)

def slidingWindow(vcf, window_size=1000):
    """Generator function that yields non overlapping chunks of VCF lines."""

    # TO DO: add overlapping increment
    chrm_id = None
    start = 0
    stop = window_size
    chunk = []

    for count, line in enumerate(vcf):

        # SKIP HEADER
        if line.startswith("#"):
            continue

        # GET CHRM AND POS
        line_parts = line.strip().split()
        chrm, pos = line_parts[:2]
        pos = int(pos)

        # UPDATE CHRM ID FOR INITAL SITE
        if chrm_id == None:
            chrm_id = chrm

        # REZERO AT NEW CHRM
        #   AND YIELD CHUNK
        if chrm_id != chrm:
            yield chunk
            stop = window_size
            chrm_id = chrm
            chunk = [line]
            continue

        # UPDATE CHUNK
        if pos < stop:
            chunk.append(line)

        # YIELD CHUNK IF CRURRENT
        #   POS EXCEEDS STOP
        if pos >= stop:
            yield chunk
            stop += window_size
            chunk = [line]

def process_window(window):
    vcf = VCF()

    Hs_est_list = []
    Ht_est_list = []

    chrm = None
    pos = None
    depth = None

    for count, line in enumerate(window):
        vcf_line_dict = vcf.parse_vcf_line(line)
        print vcf_line_dict
        if count == 0:
            chrm = vcf_line_dict["CHROM"]
            pos = vcf_line_dict['POS']
            info = vcf.parse_info(vcf_line_dict["INFO"])
            depth = info["DP"]

        allele_counts = vcf.calc_allele_counts(vcf_line_dict)
        f_statistics = vcf.calc_fstats(allele_counts)

        Hs_est = f_statistics[('MAR', 'CAP')]['Hs_est']
        Hs_est_list.append(Hs_est)

        Ht_est = f_statistics[('MAR', 'CAP')]['Ht_est']
        Ht_est_list.append(Ht_est)

    print Ht_est_list

    # # Remove uninformative SNPs (nans)
    # Hs_est_list = [item for item in Hs_est_list if not math.isnan(item)]
    # Ht_est_list = [item for item in Ht_est_list if not math.isnan(item)]

    # multilocus_f_statistics = {('MAR', 'CAP'): None}
    # if len(Hs_est_list) != 0 or len(Ht_est_list)!= 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[('MAR', 'CAP')] = values_dict

    # return vcf.__write_to_outfiles__( chrm, pos, depth, 'Gst_est', multilocus_f_statistics)

def do_windowed_analysis(vcf):

    fout = open(args.output, 'w')
    vcf_file = open(vcf.vcf_file,'rU')

    p = multiprocessing.Pool(4)

    for chunk in grouper(10, slidingWindow(vcf_file,vcf.window_size)):
        results = p.map(process_window, chunk)
        for r in results:
            print r

        # for count, window in enumerate(vcf.slidingWindow(vcf_file,args.window_size)):
        #     formated_data, header = vcf.process_window(window, args)

        #     if line_count == 0:
        #         fout.write(header+"\n")
        #         fout.write(formated_data+"\n")
        #     else:
        #         fout.write(formated_data+"\n")
        # fout.close()

if __name__ == '__main__':

    args = get_args()
    copy_reg.pickle(types.MethodType, _pickle_method, _unpickle_method)
    vcf = VCF()
    vcf.vcf_file = args.input
    vcf.set_header(vcf.vcf_file)
    vcf.populations = args.populations
    vcf.f_statistic = args.f_statistic
    vcf.window_size = args.window_size

    if args.window_size != None:
        do_windowed_analysis(vcf)

    else:
        vcf.populations = args.populations
        vcf.parse_individual_snps(args.input, args.output, args.f_statistic)

<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import unicode_literals
import tempfile

from couchexport.export import export_raw
from couchexport.models import Format
from django.test import SimpleTestCase, TestCase
from io import BytesIO

from corehq.apps.domain.shortcuts import create_domain
from corehq.apps.fixtures.exceptions import FixtureUploadError
from corehq.apps.fixtures.models import FixtureDataItem
from corehq.apps.fixtures.upload import validate_fixture_file_format
from corehq.apps.fixtures.upload.failure_messages import FAILURE_MESSAGES
from corehq.apps.fixtures.upload.run_upload import _run_fixture_upload
from corehq.apps.fixtures.upload.workbook import get_workbook
from corehq.util.files import TransientTempfile
from corehq.util.test_utils import make_make_path
from corehq.util.test_utils import generate_cases

_make_path = make_make_path(__file__)

validation_test_cases = [
    ('duplicate_tag', [
        "Lookup-tables should have unique 'table_id'. "
        "There are two rows with table_id 'things' in 'types' sheet.",

    ]),
    ("invalid_table_id", [
        "table_id 'invalid table_id' should not contain spaces or special characters, or start with a number."
    ]),
    ('multiple_errors', [
        "Excel worksheet 'level_1' does not contain the column "
        "'field: fun_fact' as specified in its 'types' definition",
        "Excel worksheet 'level_1' does not contain the column "
        "'field: other' as specified in its 'types' definition",
        "Excel worksheet 'level_2' does not contain the column "
        "'field: other' as specified in its 'types' definition",
        "There's no sheet for type 'level_3' in 'types' sheet. "
        "There must be one sheet per row in the 'types' sheet.",
    ]),
    ('type_has_no_sheet', [
        "There's no sheet for type 'things' in 'types' sheet. "
        "There must be one sheet per row in the 'types' sheet.",
    ]),
    ('has_no_field_column', [
        "Excel worksheet 'things' does not contain the column 'field: name' "
        "as specified in its 'types' definition",
    ]),
    ('has_no_field_column_extra_rows', [
        "Excel worksheet 'things' does not contain the column 'field: name' "
        "as specified in its 'types' definition",
    ]),
    ('has_extra_column', [
        "Excel worksheet 'things' has an extra column"
        "'field: fun_fact' that's not defined in its 'types' definition",
    ]),
    ('sheet_has_no_property', [
        "Excel worksheet 'things' does not contain property "
        "'lang' of the field 'name' as specified in its 'types' definition",
    ]),
    ('sheet_has_extra_property', [
        "Excel worksheet 'things' has an extra property "
        "'style' for the field 'name' that's not defined in its 'types' definition. "
        "Re-check the formatting",
    ]),
    ('invalid_field_with_property', [
        "Fields with attributes should be numbered as 'field: name integer'",
        # also triggers wrong_field_property_combos
        "Number of values for field 'name' and attribute 'lang' should be same",
    ]),
    ('invalid_property', [
        "Attribute should be written as 'name: lang integer'",
        # also triggers wrong_field_property_combos
        "Number of values for field 'name' and attribute 'lang' should be same",
    ]),
    ('wrong_field_property_combos', [
        "Number of values for field 'name' and attribute 'lang' should be same",
    ]),
    ('has_no_column', [
        "Workbook 'types' has no column 'table_id'.",
    ]),
    ('neither_fields_nor_attributes', [
        "Lookup-tables can not have empty fields and empty properties on items. "
        "table_id 'things' has no fields and no properties",
    ]),
    ('invalid_field_syntax', [
        "In Excel worksheet 'things', field 'name' should be numbered as 'field: name integer",
    ]),
    ('wrong_property_syntax', [
        "Properties should be specified as 'field 1: property 1'. In 'types' sheet, "
        "'field 1' is not correctly formatted"
    ]),
    ('invalid_field_name_numerical', [
        "Error in 'types' sheet for 'field 1', '100'. "
        "Field names should be strings, not numbers",
    ]),
    ('not_excel_file', [
        "Invalid file-format. Please upload a valid xlsx file.",
    ]),
    ('no_types_sheet', [
        "Workbook does not contain a sheet called types",
    ]),
    ('wrong_index_syntax', [
        "'field 1' is not correctly formatted in 'types' sheet. Whether a field is indexed should be specified "
        "as 'field 1: is_indexed?'. Its value should be 'yes' or 'no'.",
    ]),
    ('field_type_error', [
        "Fields with attributes should be numbered as 'field: name integer'"
    ]),
    ('property_type_error', [
        "Attribute should be written as 'name: lang integer'"
    ])
]

@generate_cases(
    validation_test_cases
)
def test_validation(self, filename, error_messages):
    upload_file = _make_path('test_upload', '{}.xlsx'.format(filename))
    if error_messages:
        with self.assertRaises(FixtureUploadError) as context:
            validate_fixture_file_format(upload_file)
        self.assertEqual(context.exception.errors, error_messages)
    else:
        # assert doesn't raise anything
        validate_fixture_file_format(upload_file)

class TestValidationComprehensiveness(SimpleTestCase):
    maxDiff = None

    def test_comprehensiveness(self):
        to_test = set(FAILURE_MESSAGES.keys())
        tested = set([validation for validation, _ in validation_test_cases])
        untested = to_test - tested
        self.assert_(
            not untested,
            "Some fixture upload errors are still untested.\n\n"
            "You have to write a test for the following fixture upload errors:\n{}"
            .format('\n'.join(untested)))

class TestFixtureUpload(TestCase):

    headers = (
        (
            'types',
            ('Delete(Y/N)', 'table_id', 'is_global?', 'field 1', 'field 1: is_indexed?')
        ),
        (
            'things',
            ('UID', 'Delete(Y/N)', 'field: name')
        )
    )

    @staticmethod
    def make_rows(item_rows):
        # given a list of fixture-items, return formatted excel rows
        return (
            (
                'types',
                ('N', 'things', 'yes', 'name', 'yes')
            ),
            (
                'things',
                tuple(item_rows)
            )
        )

    @classmethod
    def setUpClass(cls):
        super(TestFixtureUpload, cls).setUpClass()
        cls.domain = 'fixture-upload-test'
        cls.project = create_domain(cls.domain)

    @classmethod
    def tearDownClass(cls):
        cls.project.delete()
        super(TestFixtureUpload, cls).tearDownClass()

    def _get_workbook(self, filename):
        return get_workbook(_make_path('test_upload', '{}.xlsx'.format(filename)))

    def _get_workbook_from_data(self, headers, rows):
        f = BytesIO()
        export_raw(headers, rows, f, format=Format.XLS_2007)
        with TransientTempfile() as temp_path:
            with open(temp_path, 'wb') as temp_file:
                temp_file.write(f.getvalue())
            with open(temp_path, 'rb') as temp_file:
                return get_workbook(temp_file)

    def get_fixture_items(self, attribute):
        # return list of 'attribute' values of fixture table 'things'
        fixtures = []
        for fixture in FixtureDataItem.get_item_list(self.domain, 'things'):
            fixtures.append(fixture.fields.get(attribute).field_list[0].field_value)
        return fixtures

    def test_indexed_field(self):
        # test indexed fields are ready correctly
        workbook = self._get_workbook('ok')
        type_sheets = workbook.get_all_type_sheets()
        indexed_field = type_sheets[0].fields[0]
        self.assertEqual(indexed_field.field_name, 'name')
        self.assertTrue(indexed_field.is_indexed)

    def test_row_addition(self):
        # upload and then reupload with addition of a new fixture-item should create new items

        initial_rows = [(None, 'N', 'apple')]
        rows = self.make_rows(initial_rows)
        workbook = self._get_workbook_from_data(self.headers, rows)
        _run_fixture_upload(self.domain, workbook)

        self.assertListEqual(
            self.get_fixture_items('name'),
            ['apple']
        )

        # reupload with additional row
        apple_id = FixtureDataItem.get_item_list(self.domain, 'things')[0]._id
        new_rows = [(apple_id, 'N', 'apple'), (None, 'N', 'orange')]
        workbook = self._get_workbook_from_data(self.headers, self.make_rows(new_rows))
        _run_fixture_upload(self.domain, workbook)
        self.assertItemsEqual(
            self.get_fixture_items('name'),
            ['apple', 'orange']
        )
<commit_msg>Revert "fix corehq.apps.fixtures.tests.test_upload:TestFixtureUpload.test_row_addition"<commit_after>from __future__ import absolute_import
from __future__ import unicode_literals
import tempfile

from couchexport.export import export_raw
from couchexport.models import Format
from django.test import SimpleTestCase, TestCase
from io import BytesIO

from corehq.apps.domain.shortcuts import create_domain
from corehq.apps.fixtures.exceptions import FixtureUploadError
from corehq.apps.fixtures.models import FixtureDataItem
from corehq.apps.fixtures.upload import validate_fixture_file_format
from corehq.apps.fixtures.upload.failure_messages import FAILURE_MESSAGES
from corehq.apps.fixtures.upload.run_upload import _run_fixture_upload
from corehq.apps.fixtures.upload.workbook import get_workbook
from corehq.util.test_utils import make_make_path
from corehq.util.test_utils import generate_cases

_make_path = make_make_path(__file__)

validation_test_cases = [
    ('duplicate_tag', [
        "Lookup-tables should have unique 'table_id'. "
        "There are two rows with table_id 'things' in 'types' sheet.",

    ]),
    ("invalid_table_id", [
        "table_id 'invalid table_id' should not contain spaces or special characters, or start with a number."
    ]),
    ('multiple_errors', [
        "Excel worksheet 'level_1' does not contain the column "
        "'field: fun_fact' as specified in its 'types' definition",
        "Excel worksheet 'level_1' does not contain the column "
        "'field: other' as specified in its 'types' definition",
        "Excel worksheet 'level_2' does not contain the column "
        "'field: other' as specified in its 'types' definition",
        "There's no sheet for type 'level_3' in 'types' sheet. "
        "There must be one sheet per row in the 'types' sheet.",
    ]),
    ('type_has_no_sheet', [
        "There's no sheet for type 'things' in 'types' sheet. "
        "There must be one sheet per row in the 'types' sheet.",
    ]),
    ('has_no_field_column', [
        "Excel worksheet 'things' does not contain the column 'field: name' "
        "as specified in its 'types' definition",
    ]),
    ('has_no_field_column_extra_rows', [
        "Excel worksheet 'things' does not contain the column 'field: name' "
        "as specified in its 'types' definition",
    ]),
    ('has_extra_column', [
        "Excel worksheet 'things' has an extra column"
        "'field: fun_fact' that's not defined in its 'types' definition",
    ]),
    ('sheet_has_no_property', [
        "Excel worksheet 'things' does not contain property "
        "'lang' of the field 'name' as specified in its 'types' definition",
    ]),
    ('sheet_has_extra_property', [
        "Excel worksheet 'things' has an extra property "
        "'style' for the field 'name' that's not defined in its 'types' definition. "
        "Re-check the formatting",
    ]),
    ('invalid_field_with_property', [
        "Fields with attributes should be numbered as 'field: name integer'",
        # also triggers wrong_field_property_combos
        "Number of values for field 'name' and attribute 'lang' should be same",
    ]),
    ('invalid_property', [
        "Attribute should be written as 'name: lang integer'",
        # also triggers wrong_field_property_combos
        "Number of values for field 'name' and attribute 'lang' should be same",
    ]),
    ('wrong_field_property_combos', [
        "Number of values for field 'name' and attribute 'lang' should be same",
    ]),
    ('has_no_column', [
        "Workbook 'types' has no column 'table_id'.",
    ]),
    ('neither_fields_nor_attributes', [
        "Lookup-tables can not have empty fields and empty properties on items. "
        "table_id 'things' has no fields and no properties",
    ]),
    ('invalid_field_syntax', [
        "In Excel worksheet 'things', field 'name' should be numbered as 'field: name integer",
    ]),
    ('wrong_property_syntax', [
        "Properties should be specified as 'field 1: property 1'. In 'types' sheet, "
        "'field 1' is not correctly formatted"
    ]),
    ('invalid_field_name_numerical', [
        "Error in 'types' sheet for 'field 1', '100'. "
        "Field names should be strings, not numbers",
    ]),
    ('not_excel_file', [
        "Invalid file-format. Please upload a valid xlsx file.",
    ]),
    ('no_types_sheet', [
        "Workbook does not contain a sheet called types",
    ]),
    ('wrong_index_syntax', [
        "'field 1' is not correctly formatted in 'types' sheet. Whether a field is indexed should be specified "
        "as 'field 1: is_indexed?'. Its value should be 'yes' or 'no'.",
    ]),
    ('field_type_error', [
        "Fields with attributes should be numbered as 'field: name integer'"
    ]),
    ('property_type_error', [
        "Attribute should be written as 'name: lang integer'"
    ])
]

@generate_cases(
    validation_test_cases
)
def test_validation(self, filename, error_messages):
    upload_file = _make_path('test_upload', '{}.xlsx'.format(filename))
    if error_messages:
        with self.assertRaises(FixtureUploadError) as context:
            validate_fixture_file_format(upload_file)
        self.assertEqual(context.exception.errors, error_messages)
    else:
        # assert doesn't raise anything
        validate_fixture_file_format(upload_file)

class TestValidationComprehensiveness(SimpleTestCase):
    maxDiff = None

    def test_comprehensiveness(self):
        to_test = set(FAILURE_MESSAGES.keys())
        tested = set([validation for validation, _ in validation_test_cases])
        untested = to_test - tested
        self.assert_(
            not untested,
            "Some fixture upload errors are still untested.\n\n"
            "You have to write a test for the following fixture upload errors:\n{}"
            .format('\n'.join(untested)))

class TestFixtureUpload(TestCase):

    headers = (
        (
            'types',
            ('Delete(Y/N)', 'table_id', 'is_global?', 'field 1', 'field 1: is_indexed?')
        ),
        (
            'things',
            ('UID', 'Delete(Y/N)', 'field: name')
        )
    )

    @staticmethod
    def make_rows(item_rows):
        # given a list of fixture-items, return formatted excel rows
        return (
            (
                'types',
                ('N', 'things', 'yes', 'name', 'yes')
            ),
            (
                'things',
                tuple(item_rows)
            )
        )

    @classmethod
    def setUpClass(cls):
        super(TestFixtureUpload, cls).setUpClass()
        cls.domain = 'fixture-upload-test'
        cls.project = create_domain(cls.domain)

    @classmethod
    def tearDownClass(cls):
        cls.project.delete()
        super(TestFixtureUpload, cls).tearDownClass()

    def _get_workbook(self, filename):
        return get_workbook(_make_path('test_upload', '{}.xlsx'.format(filename)))

    def _get_workbook_from_data(self, headers, rows):
        file = BytesIO()
        export_raw(headers, rows, file, format=Format.XLS_2007)
        with tempfile.TemporaryFile(suffix='.xlsx') as f:
            f.write(file.getvalue())
            return get_workbook(f)

    def get_fixture_items(self, attribute):
        # return list of 'attribute' values of fixture table 'things'
        fixtures = []
        for fixture in FixtureDataItem.get_item_list(self.domain, 'things'):
            fixtures.append(fixture.fields.get(attribute).field_list[0].field_value)
        return fixtures

    def test_indexed_field(self):
        # test indexed fields are ready correctly
        workbook = self._get_workbook('ok')
        type_sheets = workbook.get_all_type_sheets()
        indexed_field = type_sheets[0].fields[0]
        self.assertEqual(indexed_field.field_name, 'name')
        self.assertTrue(indexed_field.is_indexed)

    def test_row_addition(self):
        # upload and then reupload with addition of a new fixture-item should create new items

        initial_rows = [(None, 'N', 'apple')]
        rows = self.make_rows(initial_rows)
        workbook = self._get_workbook_from_data(self.headers, rows)
        _run_fixture_upload(self.domain, workbook)

        self.assertListEqual(
            self.get_fixture_items('name'),
            ['apple']
        )

        # reupload with additional row
        apple_id = FixtureDataItem.get_item_list(self.domain, 'things')[0]._id
        new_rows = [(apple_id, 'N', 'apple'), (None, 'N', 'orange')]
        workbook = self._get_workbook_from_data(self.headers, self.make_rows(new_rows))
        _run_fixture_upload(self.domain, workbook)
        self.assertItemsEqual(
            self.get_fixture_items('name'),
            ['apple', 'orange']
        )
<|endoftext|>


<commit_before>from collections import Counter, defaultdict
import itertools
import csv
import matplotlib.pyplot as plt

from vader import SentimentIntensityAnalyzer
from nltk.corpus import tweets
from nltk.tokenize import wordpunct_tokenize

#tweets.fileids()
#tw_tokens = tweets.tokenized('tweets.20150430-223406.json')
texts = tweets.strings('tweets.20150430-223406.json')

keywords = {}
keywords['conservative'] = set(['cameron', 'david_cameron', 'davidcameron',
                                'dave', 'davecamm', 'osborne', 'portillo', 'pickles', 'tory', 'tories',
                                'torie', 'voteconservative', 'conservative', 'conservatives', 'bullingdon', 'telegraph'])
keywords['labour'] = set(['miliband', 'ed_miliband', 'edmiliband', 'edm',
                          'milliband', 'ed', 'uklabour', 'scottishlabour', 'labour', 'lab',
                          'edforchange', 'edforpm', 'milifandom', 'murphy'])
keywords['libdem'] = set(['clegg','libdem', 'libdems', 'dems', 'alexander'])
keywords['ukip'] = set(['farage', 'nigel_farage', 'nigel', 'askfarage',
                        'asknigelfarage', 'asknigelfar', 'ukip', 'davidcoburnukip'])
keywords['snp'] = set(['sturgeon', 'nicola_sturgeon', 'nicolasturgeon',
                       'nicola', 'salmond', 'snp', 'snpwin', 'votesnp', 'snpbecause', 'scotland',
                       'scotlands', 'scottish', 'indyref', 'independence', 'celebs4indy'])

def tweet_classify(text, keywords):
    toks = wordpunct_tokenize(text)
    toks_lower = [t.lower() for t in toks]
    labels = [k for k in keywords if keywords[k] & set(toks_lower)]
    return ' '.join(labels)

def full_labeled(texts, keywords, outfile=None):

    labeled_tweets = []
    sia = SentimentIntensityAnalyzer()

    for text in texts:
        labels = tweet_classify(text, keywords)
        sentiment = sia.polarity_scores(text)['compound']
        if len(labels) > 0:
            labeled_tweets.append([labels, text, sentiment])

    if outfile is not None:

        num = len(labeled_tweets)
        outfile = 'labeled_tweets.csv'

        with open(outfile, 'w', encoding= 'utf-8') as f:
            writer = csv.writer(f)
            writer.writerows(labeled_tweets)
            print('Wrote {} lines to {}'.format(num, outfile))
    else:
        return labeled_tweets

outfile = 'labels_scores.csv'
#full_labeled(texts, keywords, outfile=outfile)
#scores = partial_labeled(texts, keywords)

infile = open('labeled_tweets.csv', encoding='utf8')
reader = csv.reader(infile)

def regularize(labeled_tweets, filtered=True):
    regular = []
    for row in labeled_tweets:
        labels = row[0].split()
        polarity_score = float(row[2])
        if polarity_score == 0.0:
            continue
        if filtered and len(labels) > 1:
            continue
        for label in labels:
            regular.append([label, polarity_score])
    return regular

def means(party_scores):
    pos_scores = defaultdict(list)
    neg_scores = defaultdict(list)
    for row in party_scores:
        party = row[0]
        sentiment = row[1]
        if sentiment > 0:
            pos_scores[party].append(sentiment)
        else:
            neg_scores[party].append(sentiment)
    parties = sorted(pos_scores.keys())
    pos_means = []
    neg_means = []
    for party in parties:
        pos = pos_scores[party]
        pos_means.append(sum(pos) / len(pos))
        neg = neg_scores[party]
        neg_means.append(sum(neg) / len(neg))
    return (parties, pos_means, neg_means)

regular = regularize(reader)
(parties, pos_means, neg_means) = means(regular)
print(parties)
print(pos_means)
print(neg_means)

def plot_graph(scores):
    pos_scores = []
    neg_scores = []

    for l in sorted(scores):
        pos_value = scores[l]['pos']
        mean_pos = sum(pos_value) / len(pos_value)
        pos_scores.append(mean_pos)

        neg_value = scores[l]['neg']
        mean_neg = sum(neg_value) / len(neg_value)
        neg_scores.append(mean_neg)

    fig = plt.figure()
    x = len(scores.keys())
    ax = plt.subplot(111)
    ax.bar(x, mean_pos, width= .5, color='r')
    ax.bar(x, mean_neg, width= .5, color='b')

    plt.show()

#plot_graph(scores)

#def vectorize(toks, feats, binary = True):

    #vector = [0 for f in feats]
    #selected = [feat for tok in toks for feat in feats if occurs_in(feat, tok)]
    #counts = Counter(selected)
    #for i in range(len(feats)):
    #feat = feats[i]
    #if feat in counts:
        #if binary:
        #vector[i] = 1
        #else:
        #vector[i] = counts[feat]

    #return vector

#vectors = [vectorize(twl, WORDLIST) for twl in tw_tokens]
#vectors = [v for v in vectors if sum(v) > 0]
#vectors.sort()
#unique = list(vectors for vectors, _ in itertools.groupby(vectors))
#print(len(unique))
#for u in unique[:10]:
    #print(u)

<commit_msg>moved demo file<commit_after><|endoftext|>


<commit_before>import re;

#MD = markdown

list_level = 0;
block_tags_list = 0;
# markdown_replacing_rules =
# {
# 	#Blocks: DONE
# 	"h1": ["h1", r"^#[^#]+.*$",True],
# 	"h2": ["h2", r"^##[^#]+.*$",True],
# 	"h3": ["h3", r"^###[^#]+.*$",True],
# 	"h4": ["h4", r"^####[^#]+.*$",True],
# 	"h5": ["h4", r"^#####[^#]+.*$",True],
# 	"h6": ["h4", r"^######[^#]+.*$",True],
# 	"h1under1": ["h1", r"^====*$", True],
# 	"h1under2": ["h1", r"^----*$", True],

# 	"code" : ["code", r"^    ", True],

# 	#Inline, bold must go first, image must go first
# 	"bold" : ["b", r"\*\*[^*]+?\*\*", True],
# 	"boldunder":["b", r"__[^_]+?__",True],
# 	"emphasis": ["em", r"\*[^*]+?\*", True],
# 	"emphasisunder":["em", r"_[^_]+?_", True],
# 	"strikethrough": ["del",r"~~[^~]+?~~", True],

# 	"image" : ["img", r"\!\[.*?\]\(.*?\)", True],
# 	"link": ["a", r"\[.*?\]\(.*?\)", True],
# 	#List
# }

rules_order = [];
#While this line
#First through all the block elements, and get all the division of
 and others
#Second go through each inline elements
#Replace all the other symbol elements, (simple replace)
#Next line

#Step0
#Convert /r/n to /n and other stuff
#Step1
# Line based, search for matching and replace.
#Step2
# If nothing is replace, enclose each line in
 mode.
# Only two line break break the current mode.
#Step3
# ?
#Step4
# Profit!

class MDTag(object):
	def __init__(self,nameStr):
		self.nameStr = nameStr
	def action(self, origStr):
		""" str -> str, ismatch(bool), ismutiline
		Take a string and do manipulation to the string

		Return a string after the modification.
		bool, if the bool is true, caller shouldn't call others
		another bool, if the bool is true, the caller should continue call
		the method on the next string.
		"""
		raise Exception('Method not implemented!')

	def flush(self):
		pass

class MDRegexTag(MDTag):
	def __init__(self, nameStr, regexStr):
		self.regexObj = re.compile(regexStr)
		super(MDRegexTag, self).__init__(nameStr)

	def action(self, origStr):
		""" str -> str, ismatch(bool), ismutiline
		Take a string and do manipulation to the string

		Return a string after the modification.
		bool, if the bool is true, caller shouldn't call others
		another bool, if the bool is true, the caller should continue call
		the method on the next string.
		"""
		raise Exception('Method not implemented!')

#block level
class MDHeaderTag(MDRegexTag):
	def __init__(self):
		rawString = ""
		for i in range(6):
			if (i!=0):
				rawString += r"|"
			rawString += r"(^"+'#'*(i+1)+r"[^#\n]+.*$)"
		super(MDHeaderTag, self).__init__("header tag", rawString)

	def action(self, origStr):
		matchObj = self.regexObj.match(origStr)
		if matchObj is None:
			return ("", False ,False)
		else:
			return (self.process(matchObj), True , False)

	def process(self,matchObj):
		""" matchObj -> str
		Take a matchObj and do manipulation to the string
		Return a string after the modification
		"""
		count = 0
		newString = matchObj.group(0)
		for ch in newString:
			if ch == "#":
				count+=1
			else:
				break
		newString = newString.strip("#")
		newString = "<h{}>{}</h{}>".format(str(count), newString, str(count))
		return newString

class MDParaTag(MDTag):
	def __init__(self):
		self.capturingMode = False
		self.capturingStr = ""
		super(MDParaTag,self).__init__("p tag")

	def action(self,origStr):
		if self.capturingMode == False:
			if origStr.strip() != '\n':
				pass
		else:
			pass

class MDListTag(MDRegexTag):
	pass

#Need to happen before

class MDCodeTag(MDTag):
	def __init__(self):
		self.capturingMode = False
		self.capturedStr = ""
		super(MDCodeTag, self).__init__("code tag")

	#Private method
	def _numberOfStartingSpace(self, origStr):
		count = 0;
		for ch in origStr:
			if ch == ' ':
				count+=1
			else:
				break
		return count

	def action(self, origStr):
		if self.capturingMode == False:
			numberOfWhiteSpace = self._numberOfStartingSpace(origStr)
			if numberOfWhiteSpace >= 4 and origStr[numberOfWhiteSpace]!='\n':
				self.capturingMode = True
				self.capturedStr += origStr[numberOfWhiteSpace:]
				return ("", True, True)
			else:
				return (origStr, False, False)
		else:
			numberOfWhiteSpace = self._numberOfStartingSpace(origStr)
			if numberOfWhiteSpace >= 4:
				self.capturedStr += origStr[4:]
				return ("",True,True)
			elif numberOfWhiteSpace < 4 and origStr[numberOfWhiteSpace] == '\n':
				self.capturedStr += origStr[numberOfWhiteSpace:]
				return ("",True,True)
			else:
				self.capturingMode = False
				return_str = self.process(self.capturedStr)
				self.capturedStr = ""
				return (return_str,False,False)

	def flush(self):
		pass

	def process(self,capturedStr):
		return '

' + capturedStr + "

\n"

#Inline
class MDBoldTag(MDRegexTag):
	def __init__(self):
		super(MDBoldTag,self).__init__("bold tag", r"\*\*[^*]+?\*\*")

	def action(self,origStr):
		matchObj = self.regexObj.match(origStr)
		if matchObj:
			pass
		else:
			return ("", False, False)

class MDEmphasisTag(MDRegexTag):
	def __init__(self):
		super(MDEmphasisTag,self).__init__("emphasis tag", r"\*[^*]+?\*")

	def action(self,origStr):
		pass

class MDStrikeThroughTag(MDRegexTag):
	def __init__(self):
		super(MDStrikeThroughTag,self).__init__("strike through tag", r"~~[^~]+?~~")

	def action(self,origStr):
		pass

class MDImgTag(MDRegexTag):
	pass

class MDLinkTag(MDRegexTag):
	pass

#Parser
class MDParser:
	def __init__():
		self.orderList = [];
	"""
	Private:
		Convert the one linw to MDTag
	"""
	def _convertLineToMDTag(line):
		pass

	def parse(mdstring):
		pass
		return self.orderList

	def render():
		pass
		return html

class MDBlockParser:
	pass

class MDInlineParser:
	pass
# Test:
h = MDHeaderTag();
print h.action("######hahsdf###")

# b = MDBoldTag();
# print b.action("sDAFSd")

c = MDCodeTag();
test_code = """
     This is in code

###sdfasf###########################
    dsfsdaf

    safaf
      dsfs
sdfsdfa
    sdfaf
   sdsdf
    dsfasdf
    sdfad
    sfdaf

"""
return_str = ""
control = None
for line in test_code.splitlines(True):
	print("%%%%%%" + line)
	if (control == None):
		print("Control none")
		new_str,match,capture = c.action(line)
		print("$$Add str:" + new_str)
		return_str += new_str
		if match:
			print("match!Should not do others")
		else:
			print("Not match!Should do others")
		if capture:
			control = c
			print("capture!")
		else:
			control = None
			print("not capture!")
	else:
		print("Control is c")
		new_str,match,capture = c.action(line)
		return_str += new_str
		print("$$Add str:" + new_str)
		if match:
			print("continue match!Should not do others")
		else:
			print("Not match!Out")
		if capture:
			control = c
			print("continue capture!")
		else:
			control = None
			print("not capture!Out")
if control:
	new_str,match,capture = c.action("   dsf\n")
	return_str += new_str
print(return_str)
<commit_msg>Add inline code detection tag<commit_after>import re;

#MD = markdown

list_level = 0;
block_tags_list = 0;
# markdown_replacing_rules =
# {
# 	#Blocks: DONE
# 	"h1": ["h1", r"^#[^#]+.*$",True],
# 	"h2": ["h2", r"^##[^#]+.*$",True],
# 	"h3": ["h3", r"^###[^#]+.*$",True],
# 	"h4": ["h4", r"^####[^#]+.*$",True],
# 	"h5": ["h4", r"^#####[^#]+.*$",True],
# 	"h6": ["h4", r"^######[^#]+.*$",True],
# 	"h1under1": ["h1", r"^====*$", True],
# 	"h1under2": ["h1", r"^----*$", True],

# 	"code" : ["code", r"^    ", True],

# 	#Inline, bold must go first, image must go first
# 	"bold" : ["b", r"\*\*[^*]+?\*\*", True],
# 	"boldunder":["b", r"__[^_]+?__",True],
# 	"emphasis": ["em", r"\*[^*]+?\*", True],
# 	"emphasisunder":["em", r"_[^_]+?_", True],
# 	"strikethrough": ["del",r"~~[^~]+?~~", True],

# 	"image" : ["img", r"\!\[.*?\]\(.*?\)", True],
# 	"link": ["a", r"\[.*?\]\(.*?\)", True],
# 	#List
# }

rules_order = [];
#While this line
#First through all the block elements, and get all the division of
 and others
#Second go through each inline elements
#Replace all the other symbol elements, (simple replace)
#Next line

#Step0
#Convert /r/n to /n and other stuff
#Step1
# Line based, search for matching and replace.
#Step2
# If nothing is replace, enclose each line in
 mode.
# Only two line break break the current mode.
#Step3
# Postprocessings
#Step3
# ?
#Step4
# Profit!

class MDTag(object):
	def __init__(self,nameStr):
		self.nameStr = nameStr
	def action(self, origStr):
		""" str -> str, ismatch(bool), ismutiline
		Take a string and do manipulation to the string

		Return a string after the modification.
		bool, if the bool is true, caller shouldn't call others
		another bool, if the bool is true, the caller should continue call
		the method on the next string.
		"""
		raise Exception('Method not implemented!')

class MDRegexTag(MDTag):
	def __init__(self, nameStr, regexStr):
		self.regexObj = re.compile(regexStr)
		super(MDRegexTag, self).__init__(nameStr)

	def action(self, origStr):
		""" str -> str, ismatch(bool), ismutiline
		Take a string and do manipulation to the string

		Return a string after the modification.
		bool, if the bool is true, caller shouldn't call others
		another bool, if the bool is true, the caller should continue call
		the method on the next string.
		"""
		raise Exception('Method not implemented!')

#block level
class MDHeaderTag(MDRegexTag):
	def __init__(self):
		rawString = ""
		for i in range(6):
			if (i!=0):
				rawString += r"|"
			rawString += r"(^"+'#'*(i+1)+r"[^#\n]+.*$)"
		super(MDHeaderTag, self).__init__("header tag", rawString)

	def action(self, origStr):
		matchObj = self.regexObj.match(origStr)
		if matchObj is None:
			return ("", False ,False)
		else:
			return (self.process(matchObj), True , False)

	def process(self,matchObj):
		""" matchObj -> str
		Take a matchObj and do manipulation to the string
		Return a string after the modification
		"""
		count = 0
		newString = matchObj.group(0)
		for ch in newString:
			if ch == "#":
				count+=1
			else:
				break
		newString = newString.strip("#")
		newString = "<h{}>{}</h{}>".format(str(count), newString, str(count))
		return newString

class MDParaTag(MDTag):
	def __init__(self):
		self.capturingMode = False
		self.capturingStr = ""
		super(MDParaTag,self).__init__("p tag")

	def action(self,origStr):
		if self.capturingMode == False:
			if origStr.strip() != '\n':
				pass
		else:
			pass

class MDListTag(MDRegexTag):
	"""
	r" *[0-9]+\. .*"
	"""
	pass

class MDBlockquotes(MDRegexTag):
	pass

#Need to happen before

class MDCodeTag(MDTag):
	def __init__(self):
		self.capturingMode = False
		self.capturedStr = ""
		super(MDCodeTag, self).__init__("code tag")

	#Private method
	def _numberOfStartingSpace(self, origStr):
		count = 0;
		for ch in origStr:
			if ch == ' ':
				count+=1
			else:
				break
		return count

	def action(self, origStr):
		if self.capturingMode == False:
			numberOfWhiteSpace = self._numberOfStartingSpace(origStr)
			if numberOfWhiteSpace >= 4 and origStr[numberOfWhiteSpace]!='\n':
				self.capturingMode = True
				self.capturedStr += origStr[numberOfWhiteSpace:]
				return ("", True, True)
			else:
				return (origStr, False, False)
		else:
			numberOfWhiteSpace = self._numberOfStartingSpace(origStr)
			if numberOfWhiteSpace >= 4:
				self.capturedStr += origStr[4:]
				return ("",True,True)
			elif numberOfWhiteSpace < 4 and origStr[numberOfWhiteSpace] == '\n':
				self.capturedStr += origStr[numberOfWhiteSpace:]
				return ("",True,True)
			else:
				self.capturingMode = False
				return_str = self.process(self.capturedStr)
				self.capturedStr = ""
				return (return_str,False,False)

	def flush(self):
		pass

	def process(self,capturedStr):
		return '

' + capturedStr + "

\n"

#Inline
class MDBoldTag(MDRegexTag):
	def __init__(self):
		super(MDBoldTag,self).__init__("bold tag", r"\*\*[^*]+?\*\*")

	def action(self,origStr):
		matchObj = self.regexObj.match(origStr)
		if matchObj:
			pass
		else:
			return ("", False, False)

class MDEmphasisTag(MDRegexTag):
	def __init__(self):
		super(MDEmphasisTag,self).__init__("emphasis tag", r"\*[^*]+?\*")

	def action(self,origStr):
		pass

class MDStrikeThroughTag(MDRegexTag):
	def __init__(self):
		super(MDStrikeThroughTag,self).__init__("strike through tag", r"~~[^~]+?~~")

	def action(self,origStr):
		pass

class MDImgTag(MDRegexTag):
	def __init__(self):
		super(MDImgTag,self).__init__("img tag",r"\!\[.*?\]\(.*?\)")

	def action(self,origStr):
		pass

class MDLinkTag(MDRegexTag):
	def __init__(self):
		super(MDLinkTag,self).__init__("link tag",r"\[.*?\]\(.*?\)")
	def action(self,origStr):
		pass

class MDInlineCode(MDTag):
	def __init__(self):
		super(MDInlineCode,self).__init__("inline code tag")

	def action(self,origStr):
		#
		# 	i = 0
		# 	ignoretickCount = 0
		# 	tickPairs = []
		# 	while i < len(origStr):
		# 		# accept tick < tick count
		# 		# print(ignoretickCount)
		# 		if ignoretickCount > 0:
		# 			if origStr[i] == "`":
		# 				tickcount = self._numberOfConsecutiveTicks(origStr,i)
		# 				if tickcount >= ignoretickCount:
		# 					tickPairs.append((i, ignoretickCount))
		# 					i+= ignoretickCount
		# 					ignoretickCount = 0
		# 				else:
		# 					i+= tickcount
		# 			else:
		# 				i+= 1
		# 		# accept all tick
		# 		else:
		# 			if origStr[i] == "`":
		# 				tickcount = self._numberOfConsecutiveTicks(origStr,i)
		# 				if tickcount >= 1:
		# 					ignoretickCount = tickcount
		# 				tickPairs.append((i, tickcount))
		# 				i+= tickcount
		# 			else:
		# 				i+= 1
		# 	for each in range(len(tickPairs)/2):
		# 		print([tickPairs[each*2][0],tickPairs[each*2+1][0] + tickPairs[each*2+1][1]])
		# 		print(origStr[tickPairs[each*2][0]:tickPairs[each*2+1][0] + tickPairs[each*2+1][1]])

		# 	if len(tickPairs) % 2 == 0:
		# 		print("good")
		# 	else:
		# 		print("bad")
		i = 0
		tickList = []
		tickPairs = []
		while i <len(origStr):
			if origStr[i] == "`":
				tickcount = self._numberOfConsecutiveTicks(origStr,i)
				tickList.append([i,tickcount])
				i+= tickcount
			else:
				i+=1
		j = 0
		while j < len(tickList):
			deduction = 0
			found = False
			jchange = 0
			while found == False and tickList[j][1] - deduction > 0:
				curPosition = tickList[j][0] + deduction
				curTickCount = tickList[j][1] - deduction
				k = j + 1
				while k < len(tickList):
					if curTickCount <= tickList[k][1]:
						if (curTickCount < tickList[k][1]):
							tickList.insert(k+1,[tickList[k][0]+curTickCount, tickList[k][1]-curTickCount])
						tickList[k][1] = curTickCount
						tickPairs.append(((curPosition,curTickCount),tickList[k]))
						found = True
						inbetween = 0
						while inbetween < k-j+1:
							del tickList[j]
							inbetween += 1
						jchange = -1
						break
					k += 1
				deduction +=1
			j+= (1 + jchange)

		for each in range(len(tickPairs)):
			print([tickPairs[each][0][0],tickPairs[each][1][0] + tickPairs[each][1][1]])
			print(origStr[tickPairs[each][0][0]:tickPairs[each][1][0] + tickPairs[each][1][1]])

	def _numberOfConsecutiveTicks(self, origStr, startingIndex):
		j = startingIndex
		tickcount = 0
		while j < len(origStr):
			if origStr[j] == "`":
				tickcount += 1
			else:
				break
			j += 1
		return tickcount

#Parser
class MDParser:
	def __init__():
		self.orderList = [];
	"""
	Private:
		Convert the one linw to MDTag
	"""
	def _convertLineToMDTag(line):
		pass

	def parse(mdstring):
		pass
		return self.orderList

	def render():
		pass
		return html

class MDBlockParser:
	pass

class MDInlineParser:
	pass
# Test:
# h = MDHeaderTag();
# print h.action("######hahsdf###")

# b = MDBoldTag();
# print b.action("sDAFSd")

c = MDCodeTag();
test_code = """
     This is in code

###sdfasf###########################
    dsfsdaf

    safaf
      dsfs
sdfsdfa
    sdfaf
   sdsdf
    dsfasdf
    sdfad
    sfdaf

"""
# return_str = ""
# control = None
# for line in test_code.splitlines(True):
# 	print("%%%%%%" + line)
# 	if (control == None):
# 		print("Control none")
# 		new_str,match,capture = c.action(line)
# 		print("$$Add str:" + new_str)
# 		return_str += new_str
# 		if match:
# 			print("match!Should not do others")
# 		else:
# 			print("Not match!Should do others")
# 		if capture:
# 			control = c
# 			print("capture!")
# 		else:
# 			control = None
# 			print("not capture!")
# 	else:
# 		print("Control is c")
# 		new_str,match,capture = c.action(line)
# 		return_str += new_str
# 		print("$$Add str:" + new_str)
# 		if match:
# 			print("continue match!Should not do others")
# 		else:
# 			print("Not match!Out")
# 		if capture:
# 			control = c
# 			print("continue capture!")
# 		else:
# 			control = None
# 			print("not capture!Out")
# if control:
# 	new_str,match,capture = c.action("   dsf\n")
# 	return_str += new_str
# print(return_str)

inline = MDInlineCode()
inline.action("""
``ds`af``
``saf``
`
`sdf ` `
`` ` ``
``` `` ```
`d```
```sdaf`` sdasdf``` sdfa`saf``sdvsvzx`` `` sdaf`1sv`svzxc``
	""")

# inline = MDInlineCode()
# inline.action("""`d```
# ``""")
<|endoftext|>


<commit_before>import importlib
from importlib import abc
from .. import abc as testing_abc
from .. import util
from . import util as source_util
import imp
import marshal
import os
import sys
import types
import unittest

class PyLoaderMock(abc.PyLoader):

    # Globals that should be defined for all modules.
    source = (b"_ = '::'.join([__name__, __file__, __package__, "
              b"repr(__loader__)])")

    def __init__(self, data):
        """Take a dict of 'module_name: path' pairings.

        Paths should have no file extension, allowing packages to be denoted by
        ending in '__init__'.

        """
        self.module_paths = data
        self.path_to_module = {val:key for key,val in data.items()}

    def get_data(self, path):
        if path not in self.path_to_module:
            raise IOError
        return self.source

    def is_package(self, name):
        try:
            return '__init__' in self.module_paths[name]
        except KeyError:
            raise ImportError

    def source_path(self, name):
        try:
            return self.module_paths[name]
        except KeyError:
            raise ImportError

class PyPycLoaderMock(abc.PyPycLoader, PyLoaderMock):

    default_mtime = 1

    def __init__(self, source, bc={}):
        """Initialize mock.

        'bc' is a dict keyed on a module's name. The value is dict with
        possible keys of 'path', 'mtime', 'magic', and 'bc'. Except for 'path',
        each of those keys control if any part of created bytecode is to
        deviate from default values.

        """
        super().__init__(source)
        self.module_bytecode = {}
        self.path_to_bytecode = {}
        self.bytecode_to_path = {}
        for name, data in bc.items():
            self.path_to_bytecode[data['path']] = name
            self.bytecode_to_path[name] = data['path']
            magic = data.get('magic', imp.get_magic())
            mtime = importlib._w_long(data.get('mtime', self.default_mtime))
            if 'bc' in data:
                bc = data['bc']
            else:
                bc = self.compile_bc(name)
            self.module_bytecode[name] = magic + mtime + bc

    def compile_bc(self, name):
        source_path = self.module_paths.get(name, '<test>') or '<test>'
        code = compile(self.source, source_path, 'exec')
        return marshal.dumps(code)

    def source_mtime(self, name):
        if name in self.module_paths:
            return self.default_mtime
        elif name in self.module_bytecode:
            return None
        else:
            raise ImportError

    def bytecode_path(self, name):
        try:
            return self.bytecode_to_path[name]
        except KeyError:
            if name in self.module_paths:
                return None
            else:
                raise ImportError

    def write_bytecode(self, name, bytecode):
        self.module_bytecode[name] = bytecode
        return True

    def get_data(self, path):
        if path in self.path_to_module:
            return super().get_data(path)
        elif path in self.path_to_bytecode:
            name = self.path_to_bytecode[path]
            return self.module_bytecode[name]
        else:
            raise IOError

    def is_package(self, name):
        try:
            return super().is_package(name)
        except TypeError:
            return '__init__' in self.bytecode_to_path[name]

class PyLoaderTests(testing_abc.LoaderTests):

    """Tests for importlib.abc.PyLoader."""

    mocker = PyLoaderMock

    def eq_attrs(self, ob, **kwargs):
        for attr, val in kwargs.items():
            found = getattr(ob, attr)
            self.assertEqual(found, val,
                    "{} attribute: {} != {}".format(attr, found, val))

    def test_module(self):
        name = '<module>'
        path = 'path/to/module'
        mock = self.mocker({name: path})
        with util.uncache(name):
            module = mock.load_module(name)
            self.assert_(name in sys.modules)
        self.eq_attrs(module, __name__=name, __file__=path, __package__='',
                        __loader__=mock)
        self.assert_(not hasattr(module, '__path__'))
        return mock, name

    def test_package(self):
        name = '<pkg>'
        path = '/path/to/{}/__init__'.format(name)
        mock = self.mocker({name: path})
        with util.uncache(name):
            module = mock.load_module(name)
            self.assert_(name in sys.modules)
        self.eq_attrs(module, __name__=name, __file__=path,
                __path__=[os.path.dirname(path)], __package__=name,
                __loader__=mock)
        return mock, name

    def test_lacking_parent(self):
        name = 'pkg.mod'
        path = 'path/to/pkg/mod'
        mock = self.mocker({name: path})
        with util.uncache(name):
            module = mock.load_module(name)
            self.assert_(name in sys.modules)
        self.eq_attrs(module, __name__=name, __file__=path, __package__='pkg',
                        __loader__=mock)
        self.assert_(not hasattr(module, '__path__'))
        return mock, name

    def test_module_reuse(self):
        name = 'mod'
        path = 'path/to/mod'
        module = imp.new_module(name)
        mock = self.mocker({name: path})
        with util.uncache(name):
            sys.modules[name] = module
            loaded_module = mock.load_module(name)
            self.assert_(loaded_module is module)
            self.assert_(sys.modules[name] is module)
        return mock, name

    def test_state_after_failure(self):
        name = "mod"
        module = imp.new_module(name)
        module.blah = None
        mock = self.mocker({name: 'path/to/mod'})
        mock.source = b"1/0"
        with util.uncache(name):
            sys.modules[name] = module
            self.assertRaises(ZeroDivisionError, mock.load_module, name)
            self.assert_(sys.modules[name] is module)
            self.assert_(hasattr(module, 'blah'))
        return mock

    def test_unloadable(self):
        name = "mod"
        mock = self.mocker({name: 'path/to/mod'})
        mock.source = b"1/0"
        with util.uncache(name):
            self.assertRaises(ZeroDivisionError, mock.load_module, name)
            self.assert_(name not in sys.modules)
        return mock

class PyLoaderInterfaceTests(unittest.TestCase):

    """Tests for importlib.abc.PyLoader to make sure that when source_path()
    doesn't return a path everything works as expected."""

    def test_no_source_path(self):
        # No source path should lead to ImportError.
        name = 'mod'
        mock = PyLoaderMock({})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

    def test_source_path_is_None(self):
        name = 'mod'
        mock = PyLoaderMock({name: None})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

class PyLoaderGetSourceTests(unittest.TestCase):

    """Tests for importlib.abc.PyLoader.get_source()."""

    def test_default_encoding(self):
        # Should have no problems with UTF-8 text.
        name = 'mod'
        mock = PyLoaderMock({name: 'path/to/mod'})
        source = 'x = "ü"'
        mock.source = source.encode('utf-8')
        returned_source = mock.get_source(name)
        self.assertEqual(returned_source, source)

    def test_decoded_source(self):
        # Decoding should work.
        name = 'mod'
        mock = PyLoaderMock({name: 'path/to/mod'})
        source = "# coding: Latin-1\nx='ü'"
        assert source.encode('latin-1') != source.encode('utf-8')
        mock.source = source.encode('latin-1')
        returned_source = mock.get_source(name)
        self.assertEqual(returned_source, source)

class PyPycLoaderTests(PyLoaderTests):

    """Tests for importlib.abc.PyPycLoader."""

    mocker = PyPycLoaderMock

    @source_util.writes_bytecode
    def verify_bytecode(self, mock, name):
        assert name in mock.module_paths
        self.assert_(name in mock.module_bytecode)
        magic = mock.module_bytecode[name][:4]
        self.assertEqual(magic, imp.get_magic())
        mtime = importlib._r_long(mock.module_bytecode[name][4:8])
        self.assertEqual(mtime, 1)
        bc = mock.module_bytecode[name][8:]

    def test_module(self):
        mock, name = super().test_module()
        self.verify_bytecode(mock, name)

    def test_package(self):
        mock, name = super().test_package()
        self.verify_bytecode(mock, name)

    def test_lacking_parent(self):
        mock, name = super().test_lacking_parent()
        self.verify_bytecode(mock, name)

    def test_module_reuse(self):
        mock, name = super().test_module_reuse()
        self.verify_bytecode(mock, name)

    def test_state_after_failure(self):
        super().test_state_after_failure()

    def test_unloadable(self):
        super().test_unloadable()

class SkipWritingBytecodeTests(unittest.TestCase):

    """Test that bytecode is properly handled based on
    sys.dont_write_bytecode."""

    @source_util.writes_bytecode
    def run_test(self, dont_write_bytecode):
        name = 'mod'
        mock = PyPycLoaderMock({name: 'path/to/mod'})
        sys.dont_write_bytecode = dont_write_bytecode
        with util.uncache(name):
            mock.load_module(name)
        self.assert_((name in mock.module_bytecode) is not
                        dont_write_bytecode)

    def test_no_bytecode_written(self):
        self.run_test(True)

    def test_bytecode_written(self):
        self.run_test(False)

class RegeneratedBytecodeTests(unittest.TestCase):

    """Test that bytecode is regenerated as expected."""

    @source_util.writes_bytecode
    def test_different_magic(self):
        # A different magic number should lead to new bytecode.
        name = 'mod'
        bad_magic = b'\x00\x00\x00\x00'
        assert bad_magic != imp.get_magic()
        mock = PyPycLoaderMock({name: 'path/to/mod'},
                                {name: {'path': 'path/to/mod.bytecode',
                                        'magic': bad_magic}})
        with util.uncache(name):
            mock.load_module(name)
        self.assert_(name in mock.module_bytecode)
        magic = mock.module_bytecode[name][:4]
        self.assertEqual(magic, imp.get_magic())

    @source_util.writes_bytecode
    def test_old_mtime(self):
        # Bytecode with an older mtime should be regenerated.
        name = 'mod'
        old_mtime = PyPycLoaderMock.default_mtime - 1
        mock = PyPycLoaderMock({name: 'path/to/mod'},
                {name: {'path': 'path/to/mod.bytecode', 'mtime': old_mtime}})
        with util.uncache(name):
            mock.load_module(name)
        self.assert_(name in mock.module_bytecode)
        mtime = importlib._r_long(mock.module_bytecode[name][4:8])
        self.assertEqual(mtime, PyPycLoaderMock.default_mtime)

class BadBytecodeFailureTests(unittest.TestCase):

    """Test import failures when there is no source and parts of the bytecode
    is bad."""

    def test_bad_magic(self):
        # A bad magic number should lead to an ImportError.
        name = 'mod'
        bad_magic = b'\x00\x00\x00\x00'
        mock = PyPycLoaderMock({}, {name: {'path': 'path/to/mod',
                                            'magic': bad_magic}})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

    def test_bad_bytecode(self):
        # Bad code object bytecode should elad to an ImportError.
        name = 'mod'
        mock = PyPycLoaderMock({}, {name: {'path': '/path/to/mod', 'bc': b''}})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

def raise_ImportError(*args, **kwargs):
    raise ImportError

class MissingPathsTests(unittest.TestCase):

    """Test what happens when a source or bytecode path does not exist (either
    from *_path returning None or raising ImportError)."""

    def test_source_path_None(self):
        # Bytecode should be used when source_path returns None, along with
        # __file__ being set to the bytecode path.
        name = 'mod'
        bytecode_path = 'path/to/mod'
        mock = PyPycLoaderMock({name: None}, {name: {'path': bytecode_path}})
        with util.uncache(name):
            module = mock.load_module(name)
        self.assertEqual(module.__file__, bytecode_path)

    # Testing for bytecode_path returning None handled by all tests where no
    # bytecode initially exists.

    def test_all_paths_None(self):
        # If all *_path methods return None, raise ImportError.
        name = 'mod'
        mock = PyPycLoaderMock({name: None})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

    def test_source_path_ImportError(self):
        # An ImportError from source_path should trigger an ImportError.
        name = 'mod'
        mock = PyPycLoaderMock({}, {name: {'path': 'path/to/mod'}})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

    def test_bytecode_path_ImportError(self):
        # An ImportError from bytecode_path should trigger an ImportError.
        name = 'mod'
        mock = PyPycLoaderMock({name: 'path/to/mod'})
        bad_meth = types.MethodType(raise_ImportError, mock)
        mock.bytecode_path = bad_meth
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

def test_main():
    from test.support import run_unittest
    run_unittest(PyLoaderTests, PyLoaderInterfaceTests, PyLoaderGetSourceTests,
                    PyPycLoaderTests, SkipWritingBytecodeTests,
                    RegeneratedBytecodeTests, BadBytecodeFailureTests,
                    MissingPathsTests)

if __name__ == '__main__':
    test_main()
<commit_msg>importlib.test.source.test_abc_loader was making a bad assumption that all file paths used '/' as a path separator.<commit_after>import importlib
from importlib import abc
from .. import abc as testing_abc
from .. import util
from . import util as source_util
import imp
import marshal
import os
import sys
import types
import unittest

class PyLoaderMock(abc.PyLoader):

    # Globals that should be defined for all modules.
    source = (b"_ = '::'.join([__name__, __file__, __package__, "
              b"repr(__loader__)])")

    def __init__(self, data):
        """Take a dict of 'module_name: path' pairings.

        Paths should have no file extension, allowing packages to be denoted by
        ending in '__init__'.

        """
        self.module_paths = data
        self.path_to_module = {val:key for key,val in data.items()}

    def get_data(self, path):
        if path not in self.path_to_module:
            raise IOError
        return self.source

    def is_package(self, name):
        try:
            return '__init__' in self.module_paths[name]
        except KeyError:
            raise ImportError

    def source_path(self, name):
        try:
            return self.module_paths[name]
        except KeyError:
            raise ImportError

class PyPycLoaderMock(abc.PyPycLoader, PyLoaderMock):

    default_mtime = 1

    def __init__(self, source, bc={}):
        """Initialize mock.

        'bc' is a dict keyed on a module's name. The value is dict with
        possible keys of 'path', 'mtime', 'magic', and 'bc'. Except for 'path',
        each of those keys control if any part of created bytecode is to
        deviate from default values.

        """
        super().__init__(source)
        self.module_bytecode = {}
        self.path_to_bytecode = {}
        self.bytecode_to_path = {}
        for name, data in bc.items():
            self.path_to_bytecode[data['path']] = name
            self.bytecode_to_path[name] = data['path']
            magic = data.get('magic', imp.get_magic())
            mtime = importlib._w_long(data.get('mtime', self.default_mtime))
            if 'bc' in data:
                bc = data['bc']
            else:
                bc = self.compile_bc(name)
            self.module_bytecode[name] = magic + mtime + bc

    def compile_bc(self, name):
        source_path = self.module_paths.get(name, '<test>') or '<test>'
        code = compile(self.source, source_path, 'exec')
        return marshal.dumps(code)

    def source_mtime(self, name):
        if name in self.module_paths:
            return self.default_mtime
        elif name in self.module_bytecode:
            return None
        else:
            raise ImportError

    def bytecode_path(self, name):
        try:
            return self.bytecode_to_path[name]
        except KeyError:
            if name in self.module_paths:
                return None
            else:
                raise ImportError

    def write_bytecode(self, name, bytecode):
        self.module_bytecode[name] = bytecode
        return True

    def get_data(self, path):
        if path in self.path_to_module:
            return super().get_data(path)
        elif path in self.path_to_bytecode:
            name = self.path_to_bytecode[path]
            return self.module_bytecode[name]
        else:
            raise IOError

    def is_package(self, name):
        try:
            return super().is_package(name)
        except TypeError:
            return '__init__' in self.bytecode_to_path[name]

class PyLoaderTests(testing_abc.LoaderTests):

    """Tests for importlib.abc.PyLoader."""

    mocker = PyLoaderMock

    def eq_attrs(self, ob, **kwargs):
        for attr, val in kwargs.items():
            found = getattr(ob, attr)
            self.assertEqual(found, val,
                    "{} attribute: {} != {}".format(attr, found, val))

    def test_module(self):
        name = '<module>'
        path = os.path.join('', 'path', 'to', 'module')
        mock = self.mocker({name: path})
        with util.uncache(name):
            module = mock.load_module(name)
            self.assert_(name in sys.modules)
        self.eq_attrs(module, __name__=name, __file__=path, __package__='',
                        __loader__=mock)
        self.assert_(not hasattr(module, '__path__'))
        return mock, name

    def test_package(self):
        name = '<pkg>'
        path = os.path.join('path', 'to', name, '__init__')
        mock = self.mocker({name: path})
        with util.uncache(name):
            module = mock.load_module(name)
            self.assert_(name in sys.modules)
        self.eq_attrs(module, __name__=name, __file__=path,
                __path__=[os.path.dirname(path)], __package__=name,
                __loader__=mock)
        return mock, name

    def test_lacking_parent(self):
        name = 'pkg.mod'
        path = os.path.join('path', 'to', 'pkg', 'mod')
        mock = self.mocker({name: path})
        with util.uncache(name):
            module = mock.load_module(name)
            self.assert_(name in sys.modules)
        self.eq_attrs(module, __name__=name, __file__=path, __package__='pkg',
                        __loader__=mock)
        self.assert_(not hasattr(module, '__path__'))
        return mock, name

    def test_module_reuse(self):
        name = 'mod'
        path = os.path.join('path', 'to', 'mod')
        module = imp.new_module(name)
        mock = self.mocker({name: path})
        with util.uncache(name):
            sys.modules[name] = module
            loaded_module = mock.load_module(name)
            self.assert_(loaded_module is module)
            self.assert_(sys.modules[name] is module)
        return mock, name

    def test_state_after_failure(self):
        name = "mod"
        module = imp.new_module(name)
        module.blah = None
        mock = self.mocker({name: os.path.join('path', 'to', 'mod')})
        mock.source = b"1/0"
        with util.uncache(name):
            sys.modules[name] = module
            self.assertRaises(ZeroDivisionError, mock.load_module, name)
            self.assert_(sys.modules[name] is module)
            self.assert_(hasattr(module, 'blah'))
        return mock

    def test_unloadable(self):
        name = "mod"
        mock = self.mocker({name: os.path.join('path', 'to', 'mod')})
        mock.source = b"1/0"
        with util.uncache(name):
            self.assertRaises(ZeroDivisionError, mock.load_module, name)
            self.assert_(name not in sys.modules)
        return mock

class PyLoaderInterfaceTests(unittest.TestCase):

    """Tests for importlib.abc.PyLoader to make sure that when source_path()
    doesn't return a path everything works as expected."""

    def test_no_source_path(self):
        # No source path should lead to ImportError.
        name = 'mod'
        mock = PyLoaderMock({})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

    def test_source_path_is_None(self):
        name = 'mod'
        mock = PyLoaderMock({name: None})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

class PyLoaderGetSourceTests(unittest.TestCase):

    """Tests for importlib.abc.PyLoader.get_source()."""

    def test_default_encoding(self):
        # Should have no problems with UTF-8 text.
        name = 'mod'
        mock = PyLoaderMock({name: os.path.join('path', 'to', 'mod')})
        source = 'x = "ü"'
        mock.source = source.encode('utf-8')
        returned_source = mock.get_source(name)
        self.assertEqual(returned_source, source)

    def test_decoded_source(self):
        # Decoding should work.
        name = 'mod'
        mock = PyLoaderMock({name: os.path.join('path', 'to', 'mod')})
        source = "# coding: Latin-1\nx='ü'"
        assert source.encode('latin-1') != source.encode('utf-8')
        mock.source = source.encode('latin-1')
        returned_source = mock.get_source(name)
        self.assertEqual(returned_source, source)

class PyPycLoaderTests(PyLoaderTests):

    """Tests for importlib.abc.PyPycLoader."""

    mocker = PyPycLoaderMock

    @source_util.writes_bytecode
    def verify_bytecode(self, mock, name):
        assert name in mock.module_paths
        self.assert_(name in mock.module_bytecode)
        magic = mock.module_bytecode[name][:4]
        self.assertEqual(magic, imp.get_magic())
        mtime = importlib._r_long(mock.module_bytecode[name][4:8])
        self.assertEqual(mtime, 1)
        bc = mock.module_bytecode[name][8:]

    def test_module(self):
        mock, name = super().test_module()
        self.verify_bytecode(mock, name)

    def test_package(self):
        mock, name = super().test_package()
        self.verify_bytecode(mock, name)

    def test_lacking_parent(self):
        mock, name = super().test_lacking_parent()
        self.verify_bytecode(mock, name)

    def test_module_reuse(self):
        mock, name = super().test_module_reuse()
        self.verify_bytecode(mock, name)

    def test_state_after_failure(self):
        super().test_state_after_failure()

    def test_unloadable(self):
        super().test_unloadable()

class SkipWritingBytecodeTests(unittest.TestCase):

    """Test that bytecode is properly handled based on
    sys.dont_write_bytecode."""

    @source_util.writes_bytecode
    def run_test(self, dont_write_bytecode):
        name = 'mod'
        mock = PyPycLoaderMock({name: os.path.join('path', 'to', 'mod')})
        sys.dont_write_bytecode = dont_write_bytecode
        with util.uncache(name):
            mock.load_module(name)
        self.assert_((name in mock.module_bytecode) is not
                        dont_write_bytecode)

    def test_no_bytecode_written(self):
        self.run_test(True)

    def test_bytecode_written(self):
        self.run_test(False)

class RegeneratedBytecodeTests(unittest.TestCase):

    """Test that bytecode is regenerated as expected."""

    @source_util.writes_bytecode
    def test_different_magic(self):
        # A different magic number should lead to new bytecode.
        name = 'mod'
        bad_magic = b'\x00\x00\x00\x00'
        assert bad_magic != imp.get_magic()
        mock = PyPycLoaderMock({name: os.path.join('path', 'to', 'mod')},
                                {name: {'path': os.path.join('path', 'to',
                                                    'mod.bytecode'),
                                        'magic': bad_magic}})
        with util.uncache(name):
            mock.load_module(name)
        self.assert_(name in mock.module_bytecode)
        magic = mock.module_bytecode[name][:4]
        self.assertEqual(magic, imp.get_magic())

    @source_util.writes_bytecode
    def test_old_mtime(self):
        # Bytecode with an older mtime should be regenerated.
        name = 'mod'
        old_mtime = PyPycLoaderMock.default_mtime - 1
        mock = PyPycLoaderMock({name: os.path.join('path', 'to', 'mod')},
                {name: {'path': 'path/to/mod.bytecode', 'mtime': old_mtime}})
        with util.uncache(name):
            mock.load_module(name)
        self.assert_(name in mock.module_bytecode)
        mtime = importlib._r_long(mock.module_bytecode[name][4:8])
        self.assertEqual(mtime, PyPycLoaderMock.default_mtime)

class BadBytecodeFailureTests(unittest.TestCase):

    """Test import failures when there is no source and parts of the bytecode
    is bad."""

    def test_bad_magic(self):
        # A bad magic number should lead to an ImportError.
        name = 'mod'
        bad_magic = b'\x00\x00\x00\x00'
        mock = PyPycLoaderMock({}, {name: {'path': os.path.join('path', 'to',
                                                                'mod'),
                                            'magic': bad_magic}})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

    def test_bad_bytecode(self):
        # Bad code object bytecode should elad to an ImportError.
        name = 'mod'
        mock = PyPycLoaderMock({}, {name: {'path': os.path.join('path', 'to',
                                                                'mod'),
                                            'bc': b''}})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

def raise_ImportError(*args, **kwargs):
    raise ImportError

class MissingPathsTests(unittest.TestCase):

    """Test what happens when a source or bytecode path does not exist (either
    from *_path returning None or raising ImportError)."""

    def test_source_path_None(self):
        # Bytecode should be used when source_path returns None, along with
        # __file__ being set to the bytecode path.
        name = 'mod'
        bytecode_path = 'path/to/mod'
        mock = PyPycLoaderMock({name: None}, {name: {'path': bytecode_path}})
        with util.uncache(name):
            module = mock.load_module(name)
        self.assertEqual(module.__file__, bytecode_path)

    # Testing for bytecode_path returning None handled by all tests where no
    # bytecode initially exists.

    def test_all_paths_None(self):
        # If all *_path methods return None, raise ImportError.
        name = 'mod'
        mock = PyPycLoaderMock({name: None})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

    def test_source_path_ImportError(self):
        # An ImportError from source_path should trigger an ImportError.
        name = 'mod'
        mock = PyPycLoaderMock({}, {name: {'path': os.path.join('path', 'to',
                                                                'mod')}})
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

    def test_bytecode_path_ImportError(self):
        # An ImportError from bytecode_path should trigger an ImportError.
        name = 'mod'
        mock = PyPycLoaderMock({name: os.path.join('path', 'to', 'mod')})
        bad_meth = types.MethodType(raise_ImportError, mock)
        mock.bytecode_path = bad_meth
        with util.uncache(name):
            self.assertRaises(ImportError, mock.load_module, name)

def test_main():
    from test.support import run_unittest
    run_unittest(PyLoaderTests, PyLoaderInterfaceTests, PyLoaderGetSourceTests,
                    PyPycLoaderTests, SkipWritingBytecodeTests,
                    RegeneratedBytecodeTests, BadBytecodeFailureTests,
                    MissingPathsTests)

if __name__ == '__main__':
    test_main()
<|endoftext|>