<commit_before>import unittest
import numpy as np
from op_test import OpTest

class TestExp(OpTest):
    def setUp(self):
        self.op_type = "exp"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': np.exp(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestSigmoid(OpTest):
    def setUp(self):
        self.op_type = "sigmoid"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': 1 / (1 + np.exp(-self.inputs['X']))}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.008)

class TestTanh(OpTest):
    def setUp(self):
        self.op_type = "tanh"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': np.tanh(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestSqrt(OpTest):
    def setUp(self):
        self.op_type = "sqrt"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': np.sqrt(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestAbs(OpTest):
    def setUp(self):
        self.op_type = "abs"
        x = np.random.uniform(-1, 1, [4, 4]).astype("float32")
        # Because we set delta = 0.005 in caculating numeric gradient,
        # if x is too small, such as 0.002, x_neg will be -0.003
        # x_pos will be 0.007, so the numeric gradient is unaccurate.
        # we should avoid this
        x[np.abs(x) < 0.005] = 0.02
        self.inputs = {'X': x}
        self.outputs = {'Y': np.abs(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestRelu(OpTest):
    def setUp(self):
        self.op_type = "relu"
        x = np.random.uniform(-1, 1, [11, 17]).astype("float32")
        # The same reason with TestAbs
        x[np.abs(x) < 0.005] = 0.02
        self.inputs = {'X': x}
        self.outputs = {'Y': np.maximum(self.inputs['X'], 0)}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestBRelu(OpTest):
    def setUp(self):
        self.op_type = "brelu"
        x = np.random.uniform(-1, 1, [4, 4]).astype("float32")
        t_min = 1
        t_max = 4
        # The same with TestAbs
        x[np.abs(x - t_min) < 0.005] = t_min + 0.02
        x[np.abs(x - t_max) < 0.005] = t_max + 0.02

        self.inputs = {'X': x}
        self.attrs = {'t_min': t_min, 't_max': t_max}
        t = np.copy(x)
        t[t < t_min] = t_min
        t[t > t_max] = t_max
        self.outputs = {'Y': t}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.02)

class TestLeakyRelu(OpTest):
    def setUp(self):
        self.op_type = "leaky_relu"
        alpha = 0.02
        self.attrs = {'alpha': alpha}
        self.inputs = {'X': np.random.uniform(-3, 3, [4, 4]).astype("float32")}
        self.outputs = {
            'Y': np.maximum(self.inputs['X'], alpha * self.inputs['X'])
        }

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.008)

class TestSoftRelu(OpTest):
    def setUp(self):
        self.op_type = "soft_relu"
        x = np.random.uniform(-3, 3, [4, 4]).astype("float32")
        threshold = 2
        # The same reason with TestAbs
        x[np.abs(x - threshold) < 0.005] = threshold + 0.02
        x[np.abs(x + threshold) < 0.005] = -threshold + 0.02
        self.inputs = {'X': x}
        self.attrs = {'threshold': threshold}
        t = np.copy(x)
        t[t < -threshold] = -threshold
        t[t > threshold] = threshold
        self.outputs = {'Y': np.log((np.exp(t) + 1))}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.02)

class TestReciprocal(OpTest):
    def setUp(self):
        self.op_type = "reciprocal"
        self.inputs = {'X': np.random.uniform(1, 2, [11, 17]).astype("float32")}
        self.outputs = {'Y': np.reciprocal(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.01)

class TestLog(OpTest):
    def setUp(self):
        self.op_type = "log"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': np.log(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestSquare(OpTest):
    def setUp(self):
        self.op_type = "square"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': np.square(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestPow(OpTest):
    def setUp(self):
        self.op_type = "pow"
        self.inputs = {'X': np.random.uniform(1, 2, [11, 17]).astype("float32")}
        self.attrs = {'factor': 3}
        self.outputs = {'Y': np.power(self.inputs['X'], 3)}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.02)

class TestSTanh(OpTest):
    def setUp(self):
        self.op_type = "stanh"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        scale_a = 2.0 / 3.0
        scale_b = 1.7159
        self.attrs = {'scale_a': scale_a, 'scale_b': scale_b}
        self.outputs = {'Y': scale_b * np.tanh(self.inputs['X'] * scale_a)}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestSoftsign(OpTest):
    def setUp(self):
        self.op_type = "softsign"
        self.inputs = {
            'X': np.random.uniform(-1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {
            'Y': np.divide(self.inputs['X'], 1 + np.abs(self.inputs['X']))
        }

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

if __name__ == "__main__":
    unittest.main()
<commit_msg>Updated the reltive error<commit_after>import unittest
import numpy as np
from op_test import OpTest

class TestExp(OpTest):
    def setUp(self):
        self.op_type = "exp"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': np.exp(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestSigmoid(OpTest):
    def setUp(self):
        self.op_type = "sigmoid"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': 1 / (1 + np.exp(-self.inputs['X']))}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.008)

class TestTanh(OpTest):
    def setUp(self):
        self.op_type = "tanh"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': np.tanh(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestSqrt(OpTest):
    def setUp(self):
        self.op_type = "sqrt"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': np.sqrt(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestAbs(OpTest):
    def setUp(self):
        self.op_type = "abs"
        x = np.random.uniform(-1, 1, [4, 4]).astype("float32")
        # Because we set delta = 0.005 in caculating numeric gradient,
        # if x is too small, such as 0.002, x_neg will be -0.003
        # x_pos will be 0.007, so the numeric gradient is unaccurate.
        # we should avoid this
        x[np.abs(x) < 0.005] = 0.02
        self.inputs = {'X': x}
        self.outputs = {'Y': np.abs(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestRelu(OpTest):
    def setUp(self):
        self.op_type = "relu"
        x = np.random.uniform(-1, 1, [11, 17]).astype("float32")
        # The same reason with TestAbs
        x[np.abs(x) < 0.005] = 0.02
        self.inputs = {'X': x}
        self.outputs = {'Y': np.maximum(self.inputs['X'], 0)}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestBRelu(OpTest):
    def setUp(self):
        self.op_type = "brelu"
        x = np.random.uniform(-1, 1, [4, 4]).astype("float32")
        t_min = 1
        t_max = 4
        # The same with TestAbs
        x[np.abs(x - t_min) < 0.005] = t_min + 0.02
        x[np.abs(x - t_max) < 0.005] = t_max + 0.02

        self.inputs = {'X': x}
        self.attrs = {'t_min': t_min, 't_max': t_max}
        t = np.copy(x)
        t[t < t_min] = t_min
        t[t > t_max] = t_max
        self.outputs = {'Y': t}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.02)

class TestLeakyRelu(OpTest):
    def setUp(self):
        self.op_type = "leaky_relu"
        alpha = 0.02
        self.attrs = {'alpha': alpha}
        self.inputs = {'X': np.random.uniform(-3, 3, [4, 4]).astype("float32")}
        self.outputs = {
            'Y': np.maximum(self.inputs['X'], alpha * self.inputs['X'])
        }

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestSoftRelu(OpTest):
    def setUp(self):
        self.op_type = "soft_relu"
        x = np.random.uniform(-3, 3, [4, 4]).astype("float32")
        threshold = 2
        # The same reason with TestAbs
        x[np.abs(x - threshold) < 0.005] = threshold + 0.02
        x[np.abs(x + threshold) < 0.005] = -threshold + 0.02
        self.inputs = {'X': x}
        self.attrs = {'threshold': threshold}
        t = np.copy(x)
        t[t < -threshold] = -threshold
        t[t > threshold] = threshold
        self.outputs = {'Y': np.log((np.exp(t) + 1))}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.02)

class TestReciprocal(OpTest):
    def setUp(self):
        self.op_type = "reciprocal"
        self.inputs = {'X': np.random.uniform(1, 2, [11, 17]).astype("float32")}
        self.outputs = {'Y': np.reciprocal(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.01)

class TestLog(OpTest):
    def setUp(self):
        self.op_type = "log"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': np.log(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestSquare(OpTest):
    def setUp(self):
        self.op_type = "square"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {'Y': np.square(self.inputs['X'])}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestPow(OpTest):
    def setUp(self):
        self.op_type = "pow"
        self.inputs = {'X': np.random.uniform(1, 2, [11, 17]).astype("float32")}
        self.attrs = {'factor': 3}
        self.outputs = {'Y': np.power(self.inputs['X'], 3)}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.02)

class TestSTanh(OpTest):
    def setUp(self):
        self.op_type = "stanh"
        self.inputs = {
            'X': np.random.uniform(0.1, 1, [11, 17]).astype("float32")
        }
        scale_a = 2.0 / 3.0
        scale_b = 1.7159
        self.attrs = {'scale_a': scale_a, 'scale_b': scale_b}
        self.outputs = {'Y': scale_b * np.tanh(self.inputs['X'] * scale_a)}

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

class TestSoftsign(OpTest):
    def setUp(self):
        self.op_type = "softsign"
        self.inputs = {
            'X': np.random.uniform(-1, 1, [11, 17]).astype("float32")
        }
        self.outputs = {
            'Y': np.divide(self.inputs['X'], 1 + np.abs(self.inputs['X']))
        }

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], 'Y', max_relative_error=0.007)

if __name__ == "__main__":
    unittest.main()
<|endoftext|>


<commit_before>import os, re, datetime, json
from Bio import SeqIO
from Bio import Entrez
import requests

class parse(object):
    def __init__(self, **kwargs):
        if 'accessions' in kwargs:
            self.accessions = kwargs['accessions']
        self.gbdb = "nuccore"

    def entrez_email(self, email, **kwargs):
        #define email for entrez login
        if email is not None:
            self.email = email
        elif 'NCBI_EMAIL' in os.environ:
            self.email = os.environ['NCBI_EMAIL']
        else:
            raise Exception("Missing NCBI email")
        Entrez.email = self.email

    def parse(self, path, fname, ftype, email, **kwargs):
        if self.accessions is not None:
            accessions = [acc.strip() for acc in self.accessions.split(",")]
            self.entrez_email(email)
            gi = self.get_GIs(accessions)
            viruses, sequences = self.get_entrez_viruses(gi, **kwargs)
        elif ftype is not None and fname is not None:
            if ftype == 'genbank':
                viruses, sequences = self.parse_gb_file(path + fname, **kwargs)
            elif ftype == 'accession':
                accessions = self.parse_accession_file(path + fname, **kwargs)
                self.entrez_email(email)
                gi = self.get_GIs(accessions)
                viruses, sequences = self.get_entrez_viruses(gi, **kwargs)
            elif ftype == 'fasta':
                viruses, sequences = self.parse_fasta_file(path + fname, **kwargs)
            elif ftype == 'tsv':
                viruses, sequences = self.parse_tsv_file(path + fname, **kwargs)
                print("Parsed " + str(len(viruses)) + " viruses and " + str(len(sequences)) + " sequences from file " + path+fname)
        else:
            raise Exception("No input file name and type defined or accessions given")
        return (viruses, sequences)

    def fix_casing(self, v):
        '''
        force lower case on fields besides strain, title, authors
        '''
        for field in v:
            if field in ['strain', 'title', 'authors', 'accession']:
                pass
            elif v[field] is not None and isinstance(v[field], str):
                v[field] = v[field] = v[field].lower().replace(' ', '_')

    def fix_boolean(self, v):
        '''
        replace strings of "true" and "false" with proper booleans
        '''
        for field in v:
            if v[field] is not None:
                if v[field] == "true":
                    v[field] = True
                if v[field] == "false":
                    v[field] = False

    def parse_fasta_file(self, fasta, virus_fasta_fields, sequence_fasta_fields, **kwargs):
        '''
        Parse FASTA file with default header formatting
        :return: list of documents(dictionaries of attributes) to upload
        '''
        viruses = []
        sequences = []
        try:
            handle = open(fasta, 'r')
        except IOError:
            raise Exception(fasta, "not found")
        else:
            for record in SeqIO.parse(handle, "fasta"):
                content = list(map(lambda x: x.strip(), record.description.replace(">", "").split('|')))
                v = {key: content[ii] if ii < len(content) else "" for ii, key in virus_fasta_fields.items()}
                s = {key: content[ii] if ii < len(content) else "" for ii, key in sequence_fasta_fields.items()}
                s['sequence'] = str(record.seq).lower()
                v = self.add_virus_fields(v, **kwargs)
                s = self.add_sequence_fields(s, **kwargs)
                sequences.append(s)
                viruses.append(v)
            handle.close()
        return (viruses, sequences)

    def parse_tsv_file(self, tsv, **kwargs):
        '''
        Parse TSV file with default field ordering
        This is the same ordering as specified in 'fasta_fields'
        :return: list of documents(dictionaries of attributes) to upload
        '''
        import csv
        viruses = []
        try:
            os.path.isfile(tsv)
        except IOError:
            raise Exception(tsv, "not found")
        else:
            with open(tsv) as infile:
                table_reader = csv.reader(infile, delimiter="\t")
                header = {i: element for i, element in enumerate(next(table_reader), 0)}
                for row in table_reader:
                    v = {key: row[ii] if ii < len(row) else "" for ii, key in header.items()}
                    self.add_virus_fields(v, **kwargs)
                    self.fix_casing(v)
                    self.fix_boolean(v)
                    viruses.append(v)
        return viruses

    def parse_gb_file(self, gb, **kwargs):
        '''
        Parse genbank file
        :return: list of documents(dictionaries of attributes) to upload
        '''
        try:
            handle = open(gb, 'r')
        except IOError:
            raise Exception(gb, "not found")
        else:
            return self.parse_gb_entries(handle, **kwargs)

    def parse_accession_file(self, acc, **kwargs):
        '''
        Parse file for list of accession numbers to be uploaded to vdb
        :return: list of documents(dictionaries of attributes) to upload
        '''
        try:
            handle = open(acc, 'r')
        except IOError:
            raise Exception(acc, "not found")
        else:
            accessions = []
            for acc in handle:
                accessions.append(acc.strip())
            return accessions

    def add_virus_fields(self, v, host, country, **kwargs):
        '''
        add fields to the viruses defined at the command line
        '''
        if 'host' not in v and host is not None:
            if host == 'null' or host == 'none' or host == 'None':
                v['host'] = None
            else:
                v['host'] = host
        if 'country' not in v and country is not None:
            if country == 'null' or country == 'none' or country == 'None':
                v['country'] = None
            else :
                v['country'] = country
        v['virus'] = self.virus
        v['timestamp'] = self.rethink_io.get_upload_timestamp()
        v['virus_inclusion_date'] = self.rethink_io.get_upload_date()
        v['sequences'] = []
        v['number_sequences'] = 0
        return v

    def add_sequence_fields(self, v, locus, authors, source, url, public=True, **kwargs):
        '''
        add fields to the sequences defined aat the commandline
        '''
        if 'locus' not in v and locus is not None:
            if locus == 'null' or locus == 'none' or locus == 'None':
                v['locus'] = None
            else:
                v['locus'] = locus
        if 'authors' not in v and authors is not None:
            if authors == 'null' or authors == 'none' or authors == 'None':
                v['authors'] = None
            else:
                v['authors'] = authors
        if 'source' not in v and source is not None:
            if source == 'null' or source == 'none' or source == 'None':
                v['source'] = None
            else:
                v['source'] = source
        if 'url' not in v and url is not None:
            if url == 'null' or url == 'none' or url == 'None':
                v['url'] = None
            else:
                v['url'] = url
        if 'public' not in v and public is not None:
            v['public'] = public
        v['virus'] = self.virus
        v['timestamp'] = self.rethink_io.get_upload_timestamp()
        v['sequence_inclusion_date'] = self.rethink_io.get_upload_date()
        return v

    def get_GIs(self, accessions, **kwargs):
        '''
        Use entrez esearch to get genbank identifiers from accession numbers
        '''
        retmax = 10**5  # max records to retrieve at once; 10^5 is documented limit, but >2500 reproducibly throws errors
        queries = []
        giList = []

        for i in sorted(xrange(0, len(accessions), 2500)): # split accessions list into 2500-long portions
            queries.append(" ".join(accessions[i:i+2500])) # convert list to ' ' separated string

        assert sum([len(q.split()) for q in queries]) == len(accessions) # sanity check

        for q in queries:
            handle = Entrez.esearch(db=self.gbdb, term=q, retmax=retmax)    # retrieve xml of search results
            giList += Entrez.read(handle)['IdList'] # pull GI numbers from handle
        return giList

    def get_entrez_viruses(self, giList, **kwargs):
        '''
        Use entrez efetch to get genbank entries from genbank identifiers
        '''
        ## define batch size for download
        batchSize = 100

        # post NCBI query
        try:
            search_handle = Entrez.epost(db=self.gbdb, id=",".join(giList))
            search_results = Entrez.read(search_handle)
            webenv, query_key = search_results["WebEnv"], search_results["QueryKey"]
        except:
            print("ERROR: Couldn't connect with entrez, please run again")

        viruses = []
        sequences = []
        #fetch all results in batch of batchSize entries at once
        for start in range(0,len(giList),batchSize):
            #fetch entries in batch
            try:
                handle = Entrez.efetch(db=self.gbdb, rettype="gb", retstart=start, retmax=batchSize, webenv=webenv, query_key=query_key)
            except IOError:
                print("ERROR: Couldn't connect with entrez, please run again")
            else:
                result = self.parse_gb_entries(handle, **kwargs)
                viruses.extend(result[0])
                sequences.extend(result[1])
        return (viruses, sequences)

    def parse_gb_entries(self, handle, **kwargs):
        '''
        Go through genbank records to get relevant virus information
        '''
        viruses, sequences = [], []
        SeqIO_records = SeqIO.parse(handle, "genbank")
        for record in SeqIO_records:
            v = {}
            s = {}
            s['source'] = 'genbank'
            s['accession'] = re.match(r'^([^.]*)', record.id).group(0).upper()  # get everything before the '.'?
            s['sequence'] = str(record.seq).lower()
            print("Processing " + s['accession'])
            reference = record.annotations["references"][0]
            if reference.title is not None and reference.title != "Direct Submission":
                s['title'] = reference.title
            else:
                print("Couldn't find reference title for " + s['accession'])
                s['title'] = None
            if reference.authors is not None:
                first_author = re.match(r'^([^,]*)', reference.authors).group(0)
                s['authors'] = first_author + " et al"
            else:
                print("Couldn't parse authors for " + s['accession'])
                s['authors'] = None
                first_author = None
            url = "https://www.ncbi.nlm.nih.gov/nuccore/" + s['accession']
            s['url'] = self.get_doi_url(url, s['title'], first_author)

            record_features = record.features
            for feat in record_features:
                if feat.type == 'source':
                    qualifiers = feat.qualifiers
                    v['collection_date'] = self.convert_gb_date(qualifiers['collection_date'][0])
                    v['country'] = re.match(r'^([^:]*)', qualifiers['country'][0]).group(0)
                    if 'strain' in qualifiers:
                        v['strain'] = qualifiers['strain'][0]
                        s['strain'] = qualifiers['strain'][0]
                    elif 'isolate' in qualifiers:
                        v['strain'] = qualifiers['isolate'][0]
                        s['strain'] = qualifiers['isolate'][0]
                    else:
                        print("Couldn't parse strain name for " + s['accession'])
            v = self.add_virus_fields(v, **kwargs)
            s = self.add_sequence_fields(s, **kwargs)
            viruses.append(v)
            sequences.append(s)
            #self.fix_casing(v)
            #self.fix_boolean(v)
        handle.close()
        print("There were " + str(len(viruses)) + " viruses in the parsed file")
        return (viruses, sequences)

    def get_doi_url(self, url, title, author):
        '''
        Use crossref api to look for matching title and author name to link to DOI"
        '''
        if title is not None:
            num = str(2)
            response = json.loads(requests.get('http://api.crossref.org/works?query=%' + title + '%22&rows=' + num).text)
            items = response['message']['items']
            for item in items:
                if 'title' in item and item['title'][0] == title:
                    if 'author' in item and item['author'][0]['family'] == author:
                        if 'DOI' in item:
                            url = 'http://dx.doi.org/' + item['DOI']
        return url

    def convert_gb_date(self, collection_date):
        '''
        Converts calendar dates between given formats
        '''
        N_fields = len(collection_date.split('-'))
        if N_fields == 1:
            return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%Y'), '%Y-XX-XX')
        elif N_fields == 2:
            return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%b-%Y'), '%Y-%m-XX')
        elif N_fields == 3:
            if int(collection_date.split('-')[0]) < 32:
                return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%d-%b-%Y'), '%Y-%m-%d')
            else:
                return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%Y-%m-%d'), '%Y-%m-%d')

    def get_upload_date(self):
        return str(datetime.datetime.strftime(datetime.datetime.now(),'%Y-%m-%d'))
<commit_msg>Improve update citations<commit_after>import os, re, datetime, json
from Bio import SeqIO
from Bio import Entrez
import requests

class parse(object):
    def __init__(self, **kwargs):
        if 'accessions' in kwargs:
            self.accessions = kwargs['accessions']
        self.gbdb = "nuccore"

    def entrez_email(self, email, **kwargs):
        #define email for entrez login
        if email is not None:
            self.email = email
        elif 'NCBI_EMAIL' in os.environ:
            self.email = os.environ['NCBI_EMAIL']
        else:
            raise Exception("Missing NCBI email")
        Entrez.email = self.email

    def parse(self, path, fname, ftype, email, **kwargs):
        if self.accessions is not None:
            accessions = [acc.strip() for acc in self.accessions.split(",")]
            self.entrez_email(email)
            gi = self.get_GIs(accessions)
            viruses, sequences = self.get_entrez_viruses(gi, **kwargs)
        elif ftype is not None and fname is not None:
            if ftype == 'genbank':
                viruses, sequences = self.parse_gb_file(path + fname, **kwargs)
            elif ftype == 'accession':
                accessions = self.parse_accession_file(path + fname, **kwargs)
                self.entrez_email(email)
                gi = self.get_GIs(accessions)
                viruses, sequences = self.get_entrez_viruses(gi, **kwargs)
            elif ftype == 'fasta':
                viruses, sequences = self.parse_fasta_file(path + fname, **kwargs)
            elif ftype == 'tsv':
                viruses, sequences = self.parse_tsv_file(path + fname, **kwargs)
                print("Parsed " + str(len(viruses)) + " viruses and " + str(len(sequences)) + " sequences from file " + path+fname)
        else:
            raise Exception("No input file name and type defined or accessions given")
        return (viruses, sequences)

    def fix_casing(self, v):
        '''
        force lower case on fields besides strain, title, authors
        '''
        for field in v:
            if field in ['strain', 'title', 'authors', 'accession']:
                pass
            elif v[field] is not None and isinstance(v[field], str):
                v[field] = v[field] = v[field].lower().replace(' ', '_')

    def fix_boolean(self, v):
        '''
        replace strings of "true" and "false" with proper booleans
        '''
        for field in v:
            if v[field] is not None:
                if v[field] == "true":
                    v[field] = True
                if v[field] == "false":
                    v[field] = False

    def parse_fasta_file(self, fasta, virus_fasta_fields, sequence_fasta_fields, **kwargs):
        '''
        Parse FASTA file with default header formatting
        :return: list of documents(dictionaries of attributes) to upload
        '''
        viruses = []
        sequences = []
        try:
            handle = open(fasta, 'r')
        except IOError:
            raise Exception(fasta, "not found")
        else:
            for record in SeqIO.parse(handle, "fasta"):
                content = list(map(lambda x: x.strip(), record.description.replace(">", "").split('|')))
                v = {key: content[ii] if ii < len(content) else "" for ii, key in virus_fasta_fields.items()}
                s = {key: content[ii] if ii < len(content) else "" for ii, key in sequence_fasta_fields.items()}
                s['sequence'] = str(record.seq).lower()
                v = self.add_virus_fields(v, **kwargs)
                s = self.add_sequence_fields(s, **kwargs)
                sequences.append(s)
                viruses.append(v)
            handle.close()
        return (viruses, sequences)

    def parse_tsv_file(self, tsv, **kwargs):
        '''
        Parse TSV file with default field ordering
        This is the same ordering as specified in 'fasta_fields'
        :return: list of documents(dictionaries of attributes) to upload
        '''
        import csv
        viruses = []
        try:
            os.path.isfile(tsv)
        except IOError:
            raise Exception(tsv, "not found")
        else:
            with open(tsv) as infile:
                table_reader = csv.reader(infile, delimiter="\t")
                header = {i: element for i, element in enumerate(next(table_reader), 0)}
                for row in table_reader:
                    v = {key: row[ii] if ii < len(row) else "" for ii, key in header.items()}
                    self.add_virus_fields(v, **kwargs)
                    self.fix_casing(v)
                    self.fix_boolean(v)
                    viruses.append(v)
        return viruses

    def parse_gb_file(self, gb, **kwargs):
        '''
        Parse genbank file
        :return: list of documents(dictionaries of attributes) to upload
        '''
        try:
            handle = open(gb, 'r')
        except IOError:
            raise Exception(gb, "not found")
        else:
            return self.parse_gb_entries(handle, **kwargs)

    def parse_accession_file(self, acc, **kwargs):
        '''
        Parse file for list of accession numbers to be uploaded to vdb
        :return: list of documents(dictionaries of attributes) to upload
        '''
        try:
            handle = open(acc, 'r')
        except IOError:
            raise Exception(acc, "not found")
        else:
            accessions = []
            for acc in handle:
                accessions.append(acc.strip())
            return accessions

    def add_virus_fields(self, v, host, country, **kwargs):
        '''
        add fields to the viruses defined at the command line
        '''
        if 'host' not in v and host is not None:
            if host == 'null' or host == 'none' or host == 'None':
                v['host'] = None
            else:
                v['host'] = host
        if 'country' not in v and country is not None:
            if country == 'null' or country == 'none' or country == 'None':
                v['country'] = None
            else :
                v['country'] = country
        v['virus'] = self.virus
        v['timestamp'] = self.rethink_io.get_upload_timestamp()
        v['virus_inclusion_date'] = self.rethink_io.get_upload_date()
        v['sequences'] = []
        v['number_sequences'] = 0
        return v

    def add_sequence_fields(self, v, locus, authors, source, url, public=True, **kwargs):
        '''
        add fields to the sequences defined aat the commandline
        '''
        if 'locus' not in v and locus is not None:
            if locus == 'null' or locus == 'none' or locus == 'None':
                v['locus'] = None
            else:
                v['locus'] = locus
        if 'authors' not in v and authors is not None:
            if authors == 'null' or authors == 'none' or authors == 'None':
                v['authors'] = None
            else:
                v['authors'] = authors
        if 'source' not in v and source is not None:
            if source == 'null' or source == 'none' or source == 'None':
                v['source'] = None
            else:
                v['source'] = source
        if 'url' not in v and url is not None:
            if url == 'null' or url == 'none' or url == 'None':
                v['url'] = None
            else:
                v['url'] = url
        if 'public' not in v and public is not None:
            v['public'] = public
        v['virus'] = self.virus
        v['timestamp'] = self.rethink_io.get_upload_timestamp()
        v['sequence_inclusion_date'] = self.rethink_io.get_upload_date()
        return v

    def get_GIs(self, accessions, **kwargs):
        '''
        Use entrez esearch to get genbank identifiers from accession numbers
        '''
        retmax = 10**5  # max records to retrieve at once; 10^5 is documented limit, but >2500 reproducibly throws errors
        queries = []
        giList = []

        for i in sorted(xrange(0, len(accessions), 2500)): # split accessions list into 2500-long portions
            queries.append(" ".join(accessions[i:i+2500])) # convert list to ' ' separated string

        assert sum([len(q.split()) for q in queries]) == len(accessions) # sanity check

        for q in queries:
            handle = Entrez.esearch(db=self.gbdb, term=q, retmax=retmax)    # retrieve xml of search results
            giList += Entrez.read(handle)['IdList'] # pull GI numbers from handle
        return giList

    def get_entrez_viruses(self, giList, **kwargs):
        '''
        Use entrez efetch to get genbank entries from genbank identifiers
        '''
        ## define batch size for download
        batchSize = 100

        # post NCBI query
        try:
            search_handle = Entrez.epost(db=self.gbdb, id=",".join(giList))
            search_results = Entrez.read(search_handle)
            webenv, query_key = search_results["WebEnv"], search_results["QueryKey"]
        except:
            print("ERROR: Couldn't connect with entrez, please run again")

        viruses = []
        sequences = []
        #fetch all results in batch of batchSize entries at once
        for start in range(0,len(giList),batchSize):
            #fetch entries in batch
            try:
                handle = Entrez.efetch(db=self.gbdb, rettype="gb", retstart=start, retmax=batchSize, webenv=webenv, query_key=query_key)
            except IOError:
                print("ERROR: Couldn't connect with entrez, please run again")
            else:
                result = self.parse_gb_entries(handle, **kwargs)
                viruses.extend(result[0])
                sequences.extend(result[1])
        return (viruses, sequences)

    def parse_gb_entries(self, handle, **kwargs):
        '''
        Go through genbank records to get relevant virus information
        '''
        viruses, sequences = [], []
        SeqIO_records = SeqIO.parse(handle, "genbank")
        for record in SeqIO_records:
            v = {}
            s = {}
            s['source'] = 'genbank'
            s['accession'] = re.match(r'^([^.]*)', record.id).group(0).upper()  # get everything before the '.'?
            s['sequence'] = str(record.seq).lower()
            print("Processing " + s['accession'])
            reference = record.annotations["references"][0]
            if reference.title is not None and reference.title != "Direct Submission":
                s['title'] = reference.title
            else:
                print("Couldn't find reference title for " + s['accession'])
                s['title'] = None
            if reference.authors is not None:
                first_author = re.match(r'^([^,]*)', reference.authors).group(0)
                s['authors'] = first_author + " et al"
            else:
                print("Couldn't parse authors for " + s['accession'])
                s['authors'] = None
                first_author = None
            s['url'] = "https://www.ncbi.nlm.nih.gov/nuccore/" + s['accession']
            #s['url'] = self.get_doi_url(url, s['title'], first_author)

            record_features = record.features
            for feat in record_features:
                if feat.type == 'source':
                    qualifiers = feat.qualifiers
                    if 'collection_date' in qualifiers:
                        v['collection_date'] = self.convert_gb_date(qualifiers['collection_date'][0])
                    if 'country' in qualifiers:
                        v['country'] = re.match(r'^([^:]*)', qualifiers['country'][0]).group(0)
                    if 'strain' in qualifiers:
                        v['strain'] = qualifiers['strain'][0]
                        s['strain'] = qualifiers['strain'][0]
                    elif 'isolate' in qualifiers:
                        v['strain'] = qualifiers['isolate'][0]
                        s['strain'] = qualifiers['isolate'][0]
                    else:
                        print("Couldn't parse strain name for " + s['accession'])
            v = self.add_virus_fields(v, **kwargs)
            s = self.add_sequence_fields(s, **kwargs)
            viruses.append(v)
            sequences.append(s)
            #self.fix_casing(v)
            #self.fix_boolean(v)
        handle.close()
        print("There were " + str(len(viruses)) + " viruses in the parsed file")
        return (viruses, sequences)

    def get_doi_url(self, url, title, author):
        '''
        Use crossref api to look for matching title and author name to link to DOI
        Depcrecated for the moment. Crossref was hanging randomly.
        '''
        if title is not None:
            num = str(2)
            print("Calling crossref for " + url)
            response = json.loads(requests.get('http://api.crossref.org/works?query=%' + title + '%22&rows=' + num).text)
            items = response['message']['items']
            for item in items:
                if 'title' in item and item['title'][0] == title:
                    if 'author' in item and item['author'][0]['family'] == author:
                        if 'DOI' in item:
                            url = 'http://dx.doi.org/' + item['DOI']
        return url

    def convert_gb_date(self, collection_date):
        '''
        Converts calendar dates between given formats
        '''
        N_fields = len(collection_date.split('-'))
        if N_fields == 1:
            return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%Y'), '%Y-XX-XX')
        elif N_fields == 2:
            return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%b-%Y'), '%Y-%m-XX')
        elif N_fields == 3:
            if int(collection_date.split('-')[0]) < 32:
                return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%d-%b-%Y'), '%Y-%m-%d')
            else:
                return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%Y-%m-%d'), '%Y-%m-%d')

    def get_upload_date(self):
        return str(datetime.datetime.strftime(datetime.datetime.now(),'%Y-%m-%d'))
<|endoftext|>


<commit_before>import io
import fire
import numpy as np

import logging
logging.basicConfig(
    format='%(asctime)s [%(process)d] [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    level=logging.INFO)

def main(source_file, target_file, s2t_file, t2s_file, output_file, min_score=0.5):
    s2t = np.loadtxt(s2t_file)
    t2s = np.loadtxt(t2s_file)
    source = io.open(source_file, 'r', encoding='utf-8').read().strip().split('\n')
    target = io.open(target_file, 'r', encoding='utf-8').read().strip().split('\n')

    assert len(source) == len(target) == s2t.shape[0] == t2s.shape[0], 'mismatch sentences length'

    with io.open(output_file, 'w', encoding='utf-8') as out:
        for i, (src, tgt) in enumerate(zip(source, target)):
            s2t_score, t2s_score = s2t[i], t2s[i]
            if min([s2t_score, t2s_score]) > min_score:
                out.write('{src} ||| {tgt}\n'.format(src=src, tgt=tgt))
            if i % 1000 == 0:
                logging.info('Processed {0} sentences'.format(i))
        logging.info('Processed {0} sentences'.format(i))

if __name__ == '__main__':
    fire.Fire(main)

<commit_msg>Update filter-sym-align.py<commit_after>import io
import fire
import numpy as np

import logging
logging.basicConfig(
    format='%(asctime)s [%(process)d] [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    level=logging.INFO)

def main(source_file, target_file, s2t_file, t2s_file, output_file, min_score=0.5):
    s2t = np.loadtxt(s2t_file)
    t2s = np.loadtxt(t2s_file)
    source = io.open(source_file, 'r', encoding='utf-8').read().strip().split('\n')
    target = io.open(target_file, 'r', encoding='utf-8').read().strip().split('\n')

    assert len(source) == len(target) == s2t.shape[0] == t2s.shape[0], 'length mismatch'

    with io.open(output_file, 'w', encoding='utf-8') as out:
        for i, (src, tgt) in enumerate(zip(source, target)):
            s2t_score, t2s_score = s2t[i], t2s[i]
            if min([s2t_score, t2s_score]) > min_score:
                out.write('{src} ||| {tgt}\n'.format(src=src, tgt=tgt))
            if i % 1000 == 0:
                logging.info('Processed {0} sentences'.format(i))
        logging.info('Processed {0} sentences'.format(i))

if __name__ == '__main__':
    fire.Fire(main)

<|endoftext|>


<commit_before># Use x, y coords for unit positions
# (97, 56) ... (104, 56)
#    ...          ...
# (97, 49) ... (104, 49)
#
# Algebraic notation for a position is:
# algebraic_pos = chr(x) + chr(y)

# Accept algebraic as input; ultimately return algebraic
# Use coordinates internally
# Refer to coordinate position as 'coord'
# Refer to algebraic position as 'a1'

_coord_to_a1 = dict(
    [((x, y), chr(x) + chr(y)) for x in xrange(97, 105) for y in xrange(49, 57)]
    )
_a1_to_coord = dict(
    [(chr(x) + chr(y), (x, y)) for x in xrange(97, 105) for y in xrange(49, 57)]
    )

# Elements of visualization
pieces = u''.join(unichr(9812 + x) for x in range(12))
pieces = u' ' + pieces[:6][::-1] + pieces[6:]
allbox = u''.join(unichr(9472 + x) for x in range(200))
box = [allbox[i] for i in (2, 0, 12, 16, 20, 24, 44, 52, 28, 36, 60)]
(vbar, hbar, ul, ur, ll, lr, nt, st, wt, et, plus) = box
h3 = hbar * 3
topline = ul + (h3 + nt) * 7 + h3 + ur
midline = wt + (h3 + plus) * 7 + h3 + et
botline = ll + (h3 + st) * 7 + h3 + lr
tpl = u' {0} ' + vbar

def _is_coord_on_board(coord):
    u"""Return True if coordinate is on the board."""
    x, y = coord
    if _coord_to_a1.get((x, y), False):
        return True
    else:
        return False

class Piece(object):
    """Parent class for chess pieces."""
    def __init__(self, coord):
        """Instantiate a piece at a coordinate position."""
        super(Piece, self).__init__()
        if isinstance(coord, str):
            self.x, self.y = _a1_to_coord[coord]
        else:
            self.x, self.y = coord

    def move(self, coord, board):
        if coord in self.possible_moves(board):
            board[coord], board[(self.x, self.y)] = self, None
            self.x, self.y = coord
        return board

class SimpleUnit(Piece):
    """Returns a SimpleUnit to test position/movement basics."""
    def __init__(self, coord, color):
        super(SimpleUnit, self).__init__(coord)
        self.color = color
        self.moves = [(0, 1)]

    def possible_moves(self, board):
        valid_moves = []
        for move in self.moves:
            dx, dy = move
            new_coord = (self.x + dx, self.y + dy)
            if _is_coord_on_board(new_coord):
                if not board[new_coord] or (board[new_coord].color != self.color):
                    valid_moves.append(new_coord)
        return valid_moves

    def __repr__(self):
        if self.color == 'black':
            return 'Sb:({},{})'.format(self.x, self.y)
        else:
            return 'Sw:({},{})'.format(self.x, self.y)

class Pawn(SimpleUnit):
    """docstring for Pawn"""
    def __init__(self, coord, color):
        super(Pawn, self).__init__(coord, color)
        if color == 'black':
            self.viz = -1
        else:
            self.viz = 1

    def __repr__(self):
        if self.color == 'black':
            return 'Pb:({},{})'.format(self.x, self.y)
        else:
            return 'Pw:({},{})'.format(self.x, self.y)

class Knight(SimpleUnit):
    """docstring for Knight"""
    def __init__(self, coord, color):
        super(Knight, self).__init__(coord, color)
        if color == 'black':
            self.viz = -2
        else:
            self.viz = 2

    def __repr__(self):
        if self.color == 'black':
            return 'Kb:({},{})'.format(self.x, self.y)
        else:
            return 'Kw:({},{})'.format(self.x, self.y)

class Bishop(SimpleUnit):
    """docstring for Bishop"""
    def __init__(self, coord, color):
        super(Bishop, self).__init__(coord, color)
        if color == 'black':
            self.viz = -3
        else:
            self.viz = 3

    def __repr__(self):
        if self.color == 'black':
            return 'Bb:({},{})'.format(self.x, self.y)
        else:
            return 'Bw:({},{})'.format(self.x, self.y)

class Rook(SimpleUnit):
    """docstring for Rook"""
    def __init__(self, coord, color):
        super(Rook, self).__init__(coord, color)
        if color == 'black':
            self.viz = -4
        else:
            self.viz = 4

    def __repr__(self):
        if self.color == 'black':
            return 'Rb:({},{})'.format(self.x, self.y)
        else:
            return 'Rw:({},{})'.format(self.x, self.y)

class Queen(SimpleUnit):
    """docstring for Queen"""
    def __init__(self, coord, color):
        super(Queen, self).__init__(coord, color)
        if color == 'black':
            self.viz = -5
        else:
            self.viz = 5

    def __repr__(self):
        if self.color == 'black':
            return 'Qb:({},{})'.format(self.x, self.y)
        else:
            return 'Qw:({},{})'.format(self.x, self.y)

class King(SimpleUnit):
    """docstring for King"""
    def __init__(self, coord, color):
        super(King, self).__init__(coord, color)
        if color == 'black':
            self.viz = -6
        else:
            self.viz = 6

    def __repr__(self):
        if self.color == 'black':
            return 'Kb:({},{})'.format(self.x, self.y)
        else:
            return 'Kw:({},{})'.format(self.x, self.y)

class Match(object):

    def __init__(self):
        super(Match, self).__init__()
        self.board = self._create_blank_board()

    def move(self, start_a1, end_a1):
        """Update board with the moved piece if move is valid.

        Accepts algebraic notation.
        """
        start_coord, end_coord = _a1_to_coord[start_a1], _a1_to_coord[end_a1]
        if not self.board[start_coord]:
            raise LookupError("No piece at that location")
        piece = self.board[start_coord]
        self.board = piece.move(end_coord, self.board)
        self.view()

    def _create_blank_board(self):
        board = dict([((x, y), None) for x in xrange(97, 105) for y in xrange(49, 57)])
        return board

    def _add_simple_units(self):
        black = [(x, y) for x in xrange(97, 105) for y in xrange(55, 57)]
        white = [(x, y) for x in xrange(97, 105) for y in xrange(49, 51)]
        for i in black:
            self.board[i] = SimpleUnit(i, 'black')
        for i in white:
            self.board[i] = SimpleUnit(i, 'white')

    def _add_starting_units(self):
        black_units = [Rook, Knight, Bishop, Queen, King, Bishop, Knight, Rook]
        white_units = black_units[::-1]
        for i, unit in enumerate(black_units):
            self.board[(97 + i, 56)] = unit((97 + i, 56), 'black')
            self.board[(97 + i, 55)] = Pawn((97 + i, 55), 'black')
        for i, unit in enumerate(white_units):
            self.board[(97 + i, 49)] = unit((97 + i, 49), 'white')
            self.board[(97 + i, 50)] = Pawn((97 + i, 50), 'white')

    def _make_square(self):
        positions = [(x, y) for y in xrange(56, 48, -1) for x in xrange(97, 105)]
        square = [[] for i in xrange(8)]
        current_line = 0
        line_count = 0
        for pos in positions:
            if line_count == 8:
                current_line += 1
                line_count = 0
            if self.board[pos]:
                square[current_line].append(self.board[pos].viz)
            else:
                square[current_line].append(0)
            line_count += 1
        return square

    def _inter(self, *args):
        """Return a unicode string with a line of the chessboard.

        args are 8 integers with the values
            0 : empty square
            1, 2, 3, 4, 5, 6: white pawn, knight, bishop, rook, queen, king
            -1, -2, -3, -4, -5, -6: same black pieces
        """
        assert len(args) == 8
        return vbar + u''.join((tpl.format(pieces[a]) for a in args))

    def _game(self, position):
        yield topline
        yield self._inter(*position[0])
        for row in position[1:]:
            yield midline
            yield self._inter(*row)
        yield botline

    def view(self):
        print "\n".join(self._game(self._make_square()))

# game.__doc__ = """Return the chessboard as a string for a given position.

#     position is a list of 8 lists or tuples of length 8 containing integers
# """

if __name__ == "__main__":
    m = Match()
    m._add_starting_units()
<commit_msg>Adjust SimpleUnit to allow moves down the board if the unit is black<commit_after># Use x, y coords for unit positions
# (97, 56) ... (104, 56)
#    ...          ...
# (97, 49) ... (104, 49)
#
# Algebraic notation for a position is:
# algebraic_pos = chr(x) + chr(y)

# Accept algebraic as input; ultimately return algebraic
# Use coordinates internally
# Refer to coordinate position as 'coord'
# Refer to algebraic position as 'a1'

_coord_to_a1 = dict(
    [((x, y), chr(x) + chr(y)) for x in xrange(97, 105) for y in xrange(49, 57)]
    )
_a1_to_coord = dict(
    [(chr(x) + chr(y), (x, y)) for x in xrange(97, 105) for y in xrange(49, 57)]
    )

# Elements of visualization
pieces = u''.join(unichr(9812 + x) for x in range(12))
pieces = u' ' + pieces[:6][::-1] + pieces[6:]
allbox = u''.join(unichr(9472 + x) for x in range(200))
box = [allbox[i] for i in (2, 0, 12, 16, 20, 24, 44, 52, 28, 36, 60)]
(vbar, hbar, ul, ur, ll, lr, nt, st, wt, et, plus) = box
h3 = hbar * 3
topline = ul + (h3 + nt) * 7 + h3 + ur
midline = wt + (h3 + plus) * 7 + h3 + et
botline = ll + (h3 + st) * 7 + h3 + lr
tpl = u' {0} ' + vbar

def _is_coord_on_board(coord):
    u"""Return True if coordinate is on the board."""
    x, y = coord
    if _coord_to_a1.get((x, y), False):
        return True
    else:
        return False

class Piece(object):
    """Parent class for chess pieces."""
    def __init__(self, coord):
        """Instantiate a piece at a coordinate position."""
        super(Piece, self).__init__()
        if isinstance(coord, str):
            self.x, self.y = _a1_to_coord[coord]
        else:
            self.x, self.y = coord

    def move(self, coord, board):
        if coord in self.possible_moves(board):
            board[coord], board[(self.x, self.y)] = self, None
            self.x, self.y = coord
        return board

class SimpleUnit(Piece):
    """Returns a SimpleUnit to test position/movement basics."""
    def __init__(self, coord, color):
        super(SimpleUnit, self).__init__(coord)
        self.color = color
        if color == 'white':
            self.moves = [(0, 1)]
        else:
            self.moves = [(0, -1)]

    def possible_moves(self, board):
        valid_moves = []
        for move in self.moves:
            dx, dy = move
            new_coord = (self.x + dx, self.y + dy)
            if _is_coord_on_board(new_coord):
                if not board[new_coord] or (board[new_coord].color != self.color):
                    valid_moves.append(new_coord)
        return valid_moves

    def __repr__(self):
        if self.color == 'black':
            return 'Sb:({},{})'.format(self.x, self.y)
        else:
            return 'Sw:({},{})'.format(self.x, self.y)

class Pawn(SimpleUnit):
    """docstring for Pawn"""
    def __init__(self, coord, color):
        super(Pawn, self).__init__(coord, color)
        if color == 'black':
            self.viz = -1
        else:
            self.viz = 1

    def __repr__(self):
        if self.color == 'black':
            return 'Pb:({},{})'.format(self.x, self.y)
        else:
            return 'Pw:({},{})'.format(self.x, self.y)

class Knight(SimpleUnit):
    """docstring for Knight"""
    def __init__(self, coord, color):
        super(Knight, self).__init__(coord, color)
        if color == 'black':
            self.viz = -2
        else:
            self.viz = 2

    def __repr__(self):
        if self.color == 'black':
            return 'Kb:({},{})'.format(self.x, self.y)
        else:
            return 'Kw:({},{})'.format(self.x, self.y)

class Bishop(SimpleUnit):
    """docstring for Bishop"""
    def __init__(self, coord, color):
        super(Bishop, self).__init__(coord, color)
        if color == 'black':
            self.viz = -3
        else:
            self.viz = 3

    def __repr__(self):
        if self.color == 'black':
            return 'Bb:({},{})'.format(self.x, self.y)
        else:
            return 'Bw:({},{})'.format(self.x, self.y)

class Rook(SimpleUnit):
    """docstring for Rook"""
    def __init__(self, coord, color):
        super(Rook, self).__init__(coord, color)
        if color == 'black':
            self.viz = -4
        else:
            self.viz = 4

    def __repr__(self):
        if self.color == 'black':
            return 'Rb:({},{})'.format(self.x, self.y)
        else:
            return 'Rw:({},{})'.format(self.x, self.y)

class Queen(SimpleUnit):
    """docstring for Queen"""
    def __init__(self, coord, color):
        super(Queen, self).__init__(coord, color)
        if color == 'black':
            self.viz = -5
        else:
            self.viz = 5

    def __repr__(self):
        if self.color == 'black':
            return 'Qb:({},{})'.format(self.x, self.y)
        else:
            return 'Qw:({},{})'.format(self.x, self.y)

class King(SimpleUnit):
    """docstring for King"""
    def __init__(self, coord, color):
        super(King, self).__init__(coord, color)
        if color == 'black':
            self.viz = -6
        else:
            self.viz = 6

    def __repr__(self):
        if self.color == 'black':
            return 'Kb:({},{})'.format(self.x, self.y)
        else:
            return 'Kw:({},{})'.format(self.x, self.y)

class Match(object):

    def __init__(self):
        super(Match, self).__init__()
        self.board = self._create_blank_board()

    def move(self, start_a1, end_a1):
        """Update board with the moved piece if move is valid.

        Accepts algebraic notation.
        """
        start_coord, end_coord = _a1_to_coord[start_a1], _a1_to_coord[end_a1]
        if not self.board[start_coord]:
            raise LookupError("No piece at that location")
        piece = self.board[start_coord]
        self.board = piece.move(end_coord, self.board)
        self.view()

    def _create_blank_board(self):
        board = dict([((x, y), None) for x in xrange(97, 105) for y in xrange(49, 57)])
        return board

    def _add_simple_units(self):
        black = [(x, y) for x in xrange(97, 105) for y in xrange(55, 57)]
        white = [(x, y) for x in xrange(97, 105) for y in xrange(49, 51)]
        for i in black:
            self.board[i] = SimpleUnit(i, 'black')
        for i in white:
            self.board[i] = SimpleUnit(i, 'white')

    def _add_starting_units(self):
        black_units = [Rook, Knight, Bishop, Queen, King, Bishop, Knight, Rook]
        white_units = black_units[::-1]
        for i, unit in enumerate(black_units):
            self.board[(97 + i, 56)] = unit((97 + i, 56), 'black')
            self.board[(97 + i, 55)] = Pawn((97 + i, 55), 'black')
        for i, unit in enumerate(white_units):
            self.board[(97 + i, 49)] = unit((97 + i, 49), 'white')
            self.board[(97 + i, 50)] = Pawn((97 + i, 50), 'white')

    def _make_square(self):
        positions = [(x, y) for y in xrange(56, 48, -1) for x in xrange(97, 105)]
        square = [[] for i in xrange(8)]
        current_line = 0
        line_count = 0
        for pos in positions:
            if line_count == 8:
                current_line += 1
                line_count = 0
            if self.board[pos]:
                square[current_line].append(self.board[pos].viz)
            else:
                square[current_line].append(0)
            line_count += 1
        return square

    def _inter(self, *args):
        """Return a unicode string with a line of the chessboard.

        args are 8 integers with the values
            0 : empty square
            1, 2, 3, 4, 5, 6: white pawn, knight, bishop, rook, queen, king
            -1, -2, -3, -4, -5, -6: same black pieces
        """
        assert len(args) == 8
        return vbar + u''.join((tpl.format(pieces[a]) for a in args))

    def _game(self, position):
        yield topline
        yield self._inter(*position[0])
        for row in position[1:]:
            yield midline
            yield self._inter(*row)
        yield botline

    def view(self):
        print "\n".join(self._game(self._make_square()))

# game.__doc__ = """Return the chessboard as a string for a given position.

#     position is a list of 8 lists or tuples of length 8 containing integers
# """

if __name__ == "__main__":
    m = Match()
    m._add_starting_units()
<|endoftext|>


<commit_before>import json, sys, re, hashlib, smtplib, base64, urllib, os, csv, collections

from auth import *

from core.handler import DataHubHandler
from datahub import DataHub

from django.http import *
from django.shortcuts import render_to_response
from django.views.decorators.csrf import csrf_exempt
from django.core.context_processors import csrf
from django.core.validators import email_re
from django.db.utils import IntegrityError
from django.utils.http import urlquote_plus

from thrift.protocol.TJSONProtocol import TJSONProtocol
from thrift.transport import TTransport
from thrift.transport.TTransport import TMemoryBuffer

'''
@author: Anant Bhardwaj
@date: Mar 21, 2013

Datahub Web Handler
'''

handler = DataHubHandler()
processor = DataHub.Processor(handler)

def home(request):
  try:
    login = get_login(request)
    if login:
      return HttpResponseRedirect('/browse/%s' %(login))
    else:
      return render_to_response("home.html")
  except Exception, e:
    return HttpResponse(
        json.dumps(
          {'error': str(e)}),
        mimetype="application/json")

@login_required
def console(request):
  return render_to_response("console.html", {
    'login': get_login(request)})

@login_required
def visualize(request):
  return render_to_response("visualize.html", {
    'login': get_login(request)})

@login_required
def newrepo(request):
  return render_to_response("newrepo.html", {
    'login': get_login(request)})

@csrf_exempt
def service(request):
  try:
    iprot = TJSONProtocol(
        TMemoryBuffer(request.body))
    oprot = TJSONProtocol(TMemoryBuffer())
    processor.process(iprot, oprot)
    resp = HttpResponse(
        oprot.trans.getvalue(),
        mimetype="application/json")
    resp['Access-Control-Allow-Origin'] = "*"
    return resp
  except Exception, e:
    print str(e)
    return HttpResponse(
        json.dumps(
          {'error': str(e)}),
        mimetype="application/json")

def user(request, username):
  try:
    if(username):
      res = manager.list_repos(username)
      repos = [t[0] for t in res['tuples']]
      return render_to_response("user.html", {
          'login': get_login(request),
          'username': username,
          'repos': repos})
  except Exception, e:
    return HttpResponse(
        json.dumps(
          {'error': str(e)}),
        mimetype="application/json")

def repo(request, username, repo):
  try:
    res = manager.list_tables(username, repo)
    tables = [t[0] for t in res['tuples']]
    res = {
        'login': get_login(request),
        'username': username,
        'repo': repo,
        'tables': tables}
    res.update(csrf(request))
    return render_to_response("repo.html", res)
  except Exception, e:
    return HttpResponse(
        json.dumps(
          {'error': str(e)}),
        mimetype="application/json")

def table(request, username, repo, table):
  try:
    res = manager.execute_sql(
        username=username,
        query='SELECT * from %s.%s.%s' %(username, repo, table))
    column_names = res['column_names']
    tuples = res['tuples']
    return render_to_response("table.html", {
        'login': get_login(request),
        'username': username,
        'repo': repo,
        'table': table,
        'column_names': column_names,
        'tuples': tuples})
  except Exception, e:
    return HttpResponse(
        json.dumps(
          {'error': str(e)}),
        mimetype="application/json")

def handle_uploaded_file(username, repo, table_name, file_data):
  file_name, file_extension = os.path.splitext(str(file_data))
  with open('%s/%s/%s.%s' %(username, repo, table_name, file_extension), 'wb+') as destination:
    for chunk in f.chunks():
      destination.write(chunk)

def create_table_from_file(request):
  login = get_login(request)
  repo = ''
  if request.method == 'POST':
    data_file = request.FILES['data_file']
    file_extension = os.path.splitext(str(data_file))[1]
    data = csv.DictReader(data_file)
    table_name = request.POST['table_name']
    repo = request.POST['repo']
    columns = map(lambda x: re.sub(r'\W+', '_', x), data.fieldnames)
    dh_table_name = '%s.%s.%s' %(login, repo, table_name)
    query = 'CREATE TABLE %s (%s text' % (dh_table_name, columns[0])
    for i in range(1, len(columns)):
      query += ', %s %s' %(columns[i], 'text')

    query += ')'
    try:
      res = manager.execute_sql(username=login, query=query)
    except Exception, e:
      pass

    Tuples = collections.namedtuple('Tuples', data.fieldnames)
    tuples = [Tuples(**row) for row in data]

    for t in tuples:
      try:
        query = "INSERT INTO %s (%s) values (%s)" %(dh_table_name, ', '.join(columns), ', '.join(map(lambda x: "'" + x + "'", list(t))))
        manager.execute_sql(username=login, query=query)
      except Exception, e:
        pass

  return HttpResponseRedirect('/browse/%s/%s' %(login, repo))

<commit_msg>csv make it better<commit_after>import json, sys, re, hashlib, smtplib, base64, urllib, os, csv, collections

from auth import *

from core.handler import DataHubHandler
from datahub import DataHub

from django.http import *
from django.shortcuts import render_to_response
from django.views.decorators.csrf import csrf_exempt
from django.core.context_processors import csrf
from django.core.validators import email_re
from django.db.utils import IntegrityError
from django.utils.http import urlquote_plus

from thrift.protocol.TJSONProtocol import TJSONProtocol
from thrift.transport import TTransport
from thrift.transport.TTransport import TMemoryBuffer

'''
@author: Anant Bhardwaj
@date: Mar 21, 2013

Datahub Web Handler
'''

handler = DataHubHandler()
processor = DataHub.Processor(handler)

def home(request):
  try:
    login = get_login(request)
    if login:
      return HttpResponseRedirect('/browse/%s' %(login))
    else:
      return render_to_response("home.html")
  except Exception, e:
    return HttpResponse(
        json.dumps(
          {'error': str(e)}),
        mimetype="application/json")

@login_required
def console(request):
  return render_to_response("console.html", {
    'login': get_login(request)})

@login_required
def visualize(request):
  return render_to_response("visualize.html", {
    'login': get_login(request)})

@login_required
def newrepo(request):
  return render_to_response("newrepo.html", {
    'login': get_login(request)})

@csrf_exempt
def service(request):
  try:
    iprot = TJSONProtocol(
        TMemoryBuffer(request.body))
    oprot = TJSONProtocol(TMemoryBuffer())
    processor.process(iprot, oprot)
    resp = HttpResponse(
        oprot.trans.getvalue(),
        mimetype="application/json")
    resp['Access-Control-Allow-Origin'] = "*"
    return resp
  except Exception, e:
    print str(e)
    return HttpResponse(
        json.dumps(
          {'error': str(e)}),
        mimetype="application/json")

def user(request, username):
  try:
    if(username):
      res = manager.list_repos(username)
      repos = [t[0] for t in res['tuples']]
      return render_to_response("user.html", {
          'login': get_login(request),
          'username': username,
          'repos': repos})
  except Exception, e:
    return HttpResponse(
        json.dumps(
          {'error': str(e)}),
        mimetype="application/json")

def repo(request, username, repo):
  try:
    res = manager.list_tables(username, repo)
    tables = [t[0] for t in res['tuples']]
    res = {
        'login': get_login(request),
        'username': username,
        'repo': repo,
        'tables': tables}
    res.update(csrf(request))
    return render_to_response("repo.html", res)
  except Exception, e:
    return HttpResponse(
        json.dumps(
          {'error': str(e)}),
        mimetype="application/json")

def table(request, username, repo, table):
  try:
    res = manager.execute_sql(
        username=username,
        query='SELECT * from %s.%s.%s' %(username, repo, table))
    column_names = res['column_names']
    tuples = res['tuples']
    return render_to_response("table.html", {
        'login': get_login(request),
        'username': username,
        'repo': repo,
        'table': table,
        'column_names': column_names,
        'tuples': tuples})
  except Exception, e:
    return HttpResponse(
        json.dumps(
          {'error': str(e)}),
        mimetype="application/json")

def handle_uploaded_file(username, repo, table_name, file_data):
  file_name, file_extension = os.path.splitext(str(file_data))
  with open('%s/%s/%s.%s' %(username, repo, table_name, file_extension), 'wb+') as destination:
    for chunk in f.chunks():
      destination.write(chunk)

def create_table_from_file(request):
  login = get_login(request)
  repo = ''
  if request.method == 'POST':
    data_file = request.FILES['data_file']
    file_extension = os.path.splitext(str(data_file))[1]
    data = csv.DictReader(data_file)
    table_name = request.POST['table_name']
    repo = request.POST['repo']
    columns = map(lambda x: re.sub(r'\W+', '_', x), data.fieldnames)
    dh_table_name = '%s.%s.%s' %(login, repo, table_name)
    query = 'CREATE TABLE %s (%s text' % (dh_table_name, columns[0])
    for i in range(1, len(columns)):
      query += ', %s %s' %(columns[i], 'text')

    query += ')'
    try:
      res = manager.execute_sql(username=login, query=query)
    except Exception, e:
      pass

    Tuples = collections.namedtuple('Tuples', columns)
    tuples = [Tuples(**row) for row in data]

    for t in tuples:
      try:
        query = "INSERT INTO %s (%s) values (%s)" %(dh_table_name, ', '.join(columns), ', '.join(map(lambda x: "'" + x + "'", list(t))))
        manager.execute_sql(username=login, query=query)
      except Exception, e:
        pass

  return HttpResponseRedirect('/browse/%s/%s' %(login, repo))

<|endoftext|>


<commit_before>from datetime import datetime, timedelta
import json
import logging
import os
from os.path import splitext, basename, exists, dirname
import sys
import time

try:
    from ansible.plugins.callback import CallbackBase
except ImportError:
    # Support Ansible 1.9.x
    CallbackBase = object

import datadog

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logging.getLogger("requests").setLevel(logging.WARNING)
logging.getLogger("dd").setLevel(logging.WARNING)
LOGGER = logging.getLogger(__name__)

"""
Originally written by 'Jharrod LaFon'
#https://github.com/jlafon/ansible-profile/blob/master/callback_plugins/profile_tasks.py

"""

ANSIBLE_TIMER_LOG = os.environ.get('ANSIBLE_TIMER_LOG')

class Timestamp(object):
    """
    A class for capturing start, end and duration for an action.
    """
    def __init__(self):
        self.start = datetime.utcnow()
        self.end = None

    def stop(self):
        """
        Record the end time of the timed period.
        """
        self.end = datetime.utcnow()

    @property
    def duration(self):
        """
        Return the duration that this Timestamp covers.
        """
        return self.end - self.start

# This class only has a single method (which would ordinarily make it a
# candidate to be turned into a function). However, the TimingLoggers are
# instanciated once when ansible starts up, and then called for every play.
class TimingLogger(object):
    """
    Base-class for logging timing about ansible tasks and plays.
    """
    def log_play(self, playbook_name, playbook_timestamp, results):
        """
        Record the timing results of an ansible play.

        Arguments:
            playbook_name: the name of the playbook being logged.
            playbook_timestamp (Timestamp): the timestamps measuring how
                long the play took.
            results (dict(string -> Timestamp)): a dict mapping task names
                to Timestamps that measure how long each task took.
        """
        pass

class DatadogTimingLogger(TimingLogger):
    """
    Record ansible task and play timing to Datadog.

    Requires that the environment variable DATADOG_API_KEY be set in order
    to log any data.
    """
    def __init__(self):
        super(DatadogTimingLogger, self).__init__()

        self.datadog_api_key = os.getenv('DATADOG_API_KEY')
        self.datadog_api_initialized = False

        if self.datadog_api_key:
            datadog.initialize(
                api_key=self.datadog_api_key,
                app_key=None
            )
            self.datadog_api_initialized = True

    def clean_tag_value(self, value):
        """
        Remove any characters that aren't allowed in Datadog tags.

        Arguments:
            value: the string to be cleaned.
        """
        return value.replace(" | ", ".").replace(" ", "-").lower()

    def log_play(self, playbook_name, playbook_timestamp, results):
        if not self.datadog_api_initialized:
            return

        datadog_tasks_metrics = []
        for name, timestamp in results.items():
            datadog_tasks_metrics.append({
                'metric': 'edx.ansible.task_duration',
                'date_happened': time.mktime(timestamp.start.timetuple()),
                'points': timestamp.duration.total_seconds(),
                'tags': [
                    'task:{0}'.format(self.clean_tag_value(name)),
                    'playbook:{0}'.format(self.clean_tag_value(playbook_name))
                ]
            })
        try:
            datadog.api.Metric.send(datadog_tasks_metrics)
            datadog.api.Metric.send(
                metric="edx.ansible.playbook_duration",
                date_happened=time.mktime(playbook_timestamp.start.timetuple()),
                points=playbook_timestamp.duration.total_seconds(),
                tags=["playbook:{0}".format(self.clean_tag_value(playbook_name))]
            )
        except Exception:
            LOGGER.exception("Failed to log timing data to datadog")

class JsonTimingLogger(TimingLogger):
    """
    Record task and play timing to a local file in a JSON format.

    Requires that the environment variable ANSIBLE_TIMER_LOG be set in order
    to log any data. This specifies the file that timing data should be logged
    to. That variable can include strftime interpolation variables,
    which will be replaced with the start time of the play.
    """
    def log_play(self, playbook_name, playbook_timestamp, results):
        # N.B. This is intended to provide a consistent interface and message
        # format across all of Open edX tooling, so it deliberately eschews
        # standard python logging infrastructure.
        if ANSIBLE_TIMER_LOG is None:
            return

        messages = []
        for name, timestamp in results.items():
            messages.append({
                'task': name,
                'playbook': playbook_name,
                'started_at': timestamp.start.isoformat(),
                'ended_at': timestamp.end.isoformat(),
                'duration': timestamp.duration.total_seconds(),
            })

        messages.append({
            'playbook': playbook_name,
            'started_at': playbook_timestamp.start.isoformat(),
            'ended_at': playbook_timestamp.end.isoformat(),
            'duration': playbook_timestamp.duration.total_seconds(),
        })

        log_path = playbook_timestamp.start.strftime(ANSIBLE_TIMER_LOG)

        try:
            log_dir = dirname(log_path)
            if log_dir and not exists(log_dir):
                os.makedirs(log_dir)

            with open(log_path, 'a') as outfile:
                for log_message in messages:
                    json.dump(
                        log_message,
                        outfile,
                        separators=(',', ':'),
                        sort_keys=True,
                    )
                    outfile.write('\n')
        except Exception:
            LOGGER.exception("Unable to write json timing log messages")

class LoggingTimingLogger(TimingLogger):
    """
    Log timing information for the play and the top 10 tasks to stdout.
    """
    def log_play(self, playbook_name, playbook_timestamp, results):

        # Sort the tasks by their running time
        sorted_results = sorted(
            results.items(),
            key=lambda (task, timestamp): timestamp.duration,
            reverse=True
        )

        for name, timestamp in sorted_results[:10]:
            LOGGER.info(
                "{0:-<80}{1:->8}".format(
                    '{0} '.format(name),
                    ' {0:.02f}s'.format(timestamp.duration.total_seconds()),
                )
            )

        LOGGER.info(
            "\nPlaybook %s finished: %s, %d total tasks.  %s elapsed. \n",
            playbook_name,
            playbook_timestamp.end,
            len(results),
            playbook_timestamp.duration,
        )

class CallbackModule(CallbackBase):

    """
    Ansible plugin get the time of each task and total time
    to run the complete playbook
    """
    def __init__(self):
        self.stats = {}
        self.current_task = None
        self.playbook_name = None
        self.playbook_timestamp = None
        self.play = None

        self.loggers = [
            DatadogTimingLogger(),
            LoggingTimingLogger(),
            JsonTimingLogger(),
        ]

    def v2_playbook_on_play_start(self, play):
        self.play = play
        super(CallbackModule, self).v2_playbook_on_play_start(play)

    def playbook_on_play_start(self, pattern):
        """
        Record the start of a play.
        """
        self.playbook_name, _ = splitext(
            basename(self.play.get_name())
        )
        self.playbook_timestamp = Timestamp()

    def playbook_on_task_start(self, name, is_conditional):
        """
        Logs the start of each task
        """

        if self.current_task is not None:
            # Record the running time of the last executed task
            self.stats[self.current_task].stop()

        # Record the start time of the current task
        self.current_task = name
        self.stats[self.current_task] = Timestamp()

    def playbook_on_stats(self, stats):
        """
        Prints the timing of each task and total time to
        run the complete playbook
        """
        # Record the timing of the very last task, we use it here, because we
        # don't have stop task function by default
        if self.current_task is not None:
            self.stats[self.current_task].stop()

        self.playbook_timestamp.stop()

        for logger in self.loggers:
            logger.log_play(
                self.playbook_name,
                self.playbook_timestamp,
                self.stats,
            )
<commit_msg>Tasks can run more than once; keep all the timings<commit_after>import collections
from datetime import datetime, timedelta
import json
import logging
import os
from os.path import splitext, basename, exists, dirname
import sys
import time

try:
    from ansible.plugins.callback import CallbackBase
except ImportError:
    # Support Ansible 1.9.x
    CallbackBase = object

import datadog

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logging.getLogger("requests").setLevel(logging.WARNING)
logging.getLogger("dd").setLevel(logging.WARNING)
LOGGER = logging.getLogger(__name__)

"""
Originally written by 'Jharrod LaFon'
#https://github.com/jlafon/ansible-profile/blob/master/callback_plugins/profile_tasks.py

"""

ANSIBLE_TIMER_LOG = os.environ.get('ANSIBLE_TIMER_LOG')

class Timestamp(object):
    """
    A class for capturing start, end and duration for an action.
    """
    def __init__(self):
        self.start = datetime.utcnow()
        self.end = None

    def stop(self):
        """
        Record the end time of the timed period.
        """
        self.end = datetime.utcnow()

    @property
    def duration(self):
        """
        Return the duration that this Timestamp covers.
        """
        return self.end - self.start

# This class only has a single method (which would ordinarily make it a
# candidate to be turned into a function). However, the TimingLoggers are
# instanciated once when ansible starts up, and then called for every play.
class TimingLogger(object):
    """
    Base-class for logging timing about ansible tasks and plays.
    """
    def log_play(self, playbook_name, playbook_timestamp, results):
        """
        Record the timing results of an ansible play.

        Arguments:
            playbook_name: the name of the playbook being logged.
            playbook_timestamp (Timestamp): the timestamps measuring how
                long the play took.
            results (dict(string -> Timestamp)): a dict mapping task names
                to Timestamps that measure how long each task took.
        """
        pass

class DatadogTimingLogger(TimingLogger):
    """
    Record ansible task and play timing to Datadog.

    Requires that the environment variable DATADOG_API_KEY be set in order
    to log any data.
    """
    def __init__(self):
        super(DatadogTimingLogger, self).__init__()

        self.datadog_api_key = os.getenv('DATADOG_API_KEY')
        self.datadog_api_initialized = False

        if self.datadog_api_key:
            datadog.initialize(
                api_key=self.datadog_api_key,
                app_key=None
            )
            self.datadog_api_initialized = True

    def clean_tag_value(self, value):
        """
        Remove any characters that aren't allowed in Datadog tags.

        Arguments:
            value: the string to be cleaned.
        """
        return value.replace(" | ", ".").replace(" ", "-").lower()

    def log_play(self, playbook_name, playbook_timestamp, results):
        if not self.datadog_api_initialized:
            return

        datadog_tasks_metrics = []
        for name, timestamp in results.items():
            datadog_tasks_metrics.append({
                'metric': 'edx.ansible.task_duration',
                'date_happened': time.mktime(timestamp.start.timetuple()),
                'points': timestamp.duration.total_seconds(),
                'tags': [
                    'task:{0}'.format(self.clean_tag_value(name)),
                    'playbook:{0}'.format(self.clean_tag_value(playbook_name))
                ]
            })
        try:
            datadog.api.Metric.send(datadog_tasks_metrics)
            datadog.api.Metric.send(
                metric="edx.ansible.playbook_duration",
                date_happened=time.mktime(playbook_timestamp.start.timetuple()),
                points=playbook_timestamp.duration.total_seconds(),
                tags=["playbook:{0}".format(self.clean_tag_value(playbook_name))]
            )
        except Exception:
            LOGGER.exception("Failed to log timing data to datadog")

class JsonTimingLogger(TimingLogger):
    """
    Record task and play timing to a local file in a JSON format.

    Requires that the environment variable ANSIBLE_TIMER_LOG be set in order
    to log any data. This specifies the file that timing data should be logged
    to. That variable can include strftime interpolation variables,
    which will be replaced with the start time of the play.
    """
    def log_play(self, playbook_name, playbook_timestamp, results):
        # N.B. This is intended to provide a consistent interface and message
        # format across all of Open edX tooling, so it deliberately eschews
        # standard python logging infrastructure.
        if ANSIBLE_TIMER_LOG is None:
            return

        messages = []
        for name, timestamp in results.items():
            messages.append({
                'task': name,
                'playbook': playbook_name,
                'started_at': timestamp.start.isoformat(),
                'ended_at': timestamp.end.isoformat(),
                'duration': timestamp.duration.total_seconds(),
            })

        messages.append({
            'playbook': playbook_name,
            'started_at': playbook_timestamp.start.isoformat(),
            'ended_at': playbook_timestamp.end.isoformat(),
            'duration': playbook_timestamp.duration.total_seconds(),
        })

        log_path = playbook_timestamp.start.strftime(ANSIBLE_TIMER_LOG)

        try:
            log_dir = dirname(log_path)
            if log_dir and not exists(log_dir):
                os.makedirs(log_dir)

            with open(log_path, 'a') as outfile:
                for log_message in messages:
                    json.dump(
                        log_message,
                        outfile,
                        separators=(',', ':'),
                        sort_keys=True,
                    )
                    outfile.write('\n')
        except Exception:
            LOGGER.exception("Unable to write json timing log messages")

class LoggingTimingLogger(TimingLogger):
    """
    Log timing information for the play and the top 10 tasks to stdout.
    """
    def log_play(self, playbook_name, playbook_timestamp, results):

        # Sort the tasks by their running time
        sorted_results = sorted(
            results.items(),
            key=lambda (task, timestamp): timestamp.duration,
            reverse=True
        )

        for name, timestamp in sorted_results[:10]:
            LOGGER.info(
                "{0:-<80}{1:->8}".format(
                    ' {0} '.format(name),
                    ' {0:.02f}s'.format(timestamp.duration.total_seconds()),
                )
            )

        LOGGER.info(
            "\nPlaybook %s finished: %s, %d total tasks.  %s elapsed. \n",
            playbook_name,
            playbook_timestamp.end,
            len(results),
            playbook_timestamp.duration,
        )

class CallbackModule(CallbackBase):

    """
    Ansible plugin get the time of each task and total time
    to run the complete playbook
    """
    def __init__(self):
        self.stats = collections.defaultdict(list)
        self.current_task = None
        self.playbook_name = None
        self.playbook_timestamp = None
        self.play = None

        self.loggers = [
            DatadogTimingLogger(),
            LoggingTimingLogger(),
            JsonTimingLogger(),
        ]

    def v2_playbook_on_play_start(self, play):
        self.play = play
        super(CallbackModule, self).v2_playbook_on_play_start(play)

    def playbook_on_play_start(self, pattern):
        """
        Record the start of a play.
        """
        self.playbook_name, _ = splitext(
            basename(self.play.get_name())
        )
        self.playbook_timestamp = Timestamp()

    def playbook_on_task_start(self, name, is_conditional):
        """
        Logs the start of each task
        """

        if self.current_task is not None:
            # Record the running time of the last executed task
            self.stats[self.current_task][-1].stop()

        # Record the start time of the current task
        self.current_task = name
        self.stats[self.current_task].append(Timestamp())

    def playbook_on_stats(self, stats):
        """
        Prints the timing of each task and total time to
        run the complete playbook
        """
        # Record the timing of the very last task, we use it here, because we
        # don't have stop task function by default
        if self.current_task is not None:
            self.stats[self.current_task][-1].stop()

        self.playbook_timestamp.stop()

        # Flatten the stats so that multiple runs of the same task get listed
        # individually.
        flat_stats = {}
        for task_name, task_runs in self.stats.iteritems():
            if len(task_runs) == 1:
                flat_stats[task_name] = task_runs[0]
            else:
                for i, run in enumerate(task_runs):
                    run_name = "{} [{}]".format(task_name, i)
                    flat_stats[run_name] = run

        for logger in self.loggers:
            logger.log_play(
                self.playbook_name,
                self.playbook_timestamp,
                flat_stats,
            )
<|endoftext|>


<commit_before>if "bpy" in locals():
    import imp
    # imp.reload(renderer)
    # imp.reload(converter)
else:
    import bpy

DEFAULT_WATCHLIST = [
    "actions",
    "armatures",
    "cameras",
    "images",
    "lamps",
    "materials",
    "meshes",
    "objects",
    "scenes",
    "sounds",
    "speakers",
    "textures",
    "worlds",
]

def get_collection_name(collection):
    class_name = collection.rna_type.__class__.__name__
    clean_name = class_name.replace("BlendData", "").lower()
    return clean_name

class RealTimeEngine(bpy.types.RenderEngine):
    bl_idname = 'RTE_FRAMEWORK'
    bl_label = "Real Time Engine Framework"

    def __init__(self, watch_list=DEFAULT_WATCHLIST):
        self._watch_list = [getattr(bpy.data, i) for i in watch_list]

        self._tracking_sets = {}
        for collection in self._watch_list:
            collection_name = get_collection_name(collection)
            self._tracking_sets[collection_name] = set()

        self._old_vmat = None
        self._old_pmat = None
        self._old_viewport = None

        def main_loop(scene):
            try:
                self.scene_callback()
            except ReferenceError:
                bpy.apps.handlers.scene_update_post.remove(main_loop)

        bpy.app.handlers.scene_update_post.append(main_loop)

    def view_update(self, context):
        """ Called when the scene is changed """
        for collection in self._watch_list:
            collection_name = get_collection_name(collection)
            collection_set = set(collection)
            tracking_set = self._tracking_sets[collection_name]

            # Check for new items
            add_method = getattr(self, "add_"+collection_name, self.add_data)
            add_set = collection_set - tracking_set
            add_method(add_set)
            tracking_set |= add_set

            # Check for removed items
            remove_method = getattr(self, "remove_"+collection_name, self.remove_data)
            remove_set = tracking_set - collection_set
            remove_method(remove_set)
            tracking_set -= remove_set

            # Check for updates
            update_method = getattr(self, "update_"+collection_name, self.update_data)
            update_set = [item for item in collection if item.is_updated]
            update_method(update_set)

    def view_draw(self, context):
        """ Called when viewport settings change """
        region = context.region
        view = context.region_data

        vmat = view.view_matrix.copy()
        vmat_inv = vmat.inverted()
        pmat = view.perspective_matrix * vmat_inv

        viewport = [region.x, region.y, region.width, region.height]

        self.update_view(vmat, pmat, viewport)

    def add_objects(self, add_set):
        for data in add_set:
            print(data.__class__.__name__, data.name, "added with Object add")

    def remove_objects(self, add_set):
        for data in add_set:
            print(data.__class__.__name__, data.name, "removed with Object remove")

    def update_objects(self, add_set):
        for data in add_set:
            print(data.__class__.__name__, data.name, "updated with Object update")

    def add_data(self, add_set):
        for data in add_set:
            print(data.__class__.__name__, data.name, "added with generic add")

    def remove_data(self, remove_set):
        for data in remove_set:
            print(data.__class__.__name__, data.name, "removed with generic remove")

    def update_data(self, update_set):
        for data in update_set:
            print(data.__class__.__name__, data.name, "updated with generic update")

    def update_view(self, view_matrix, projection_matrix, viewport):
        print("update_view")
        if view_matrix != self._old_vmat:
            print(view_matrix)
        if projection_matrix != self._old_pmat:
            print(projection_matrix)
        if viewport != viewport:
            print(viewport)

    def scene_callback(self):
        pass

<commit_msg>Display: Rendering image data in self.display<commit_after>if "bpy" in locals():
    import imp
    # imp.reload(renderer)
    # imp.reload(converter)
else:
    import bpy

from bgl import *

DEFAULT_WATCHLIST = [
    "actions",
    "armatures",
    "cameras",
    "images",
    "lamps",
    "materials",
    "meshes",
    "objects",
    "scenes",
    "sounds",
    "speakers",
    "textures",
    "worlds",
]

def get_collection_name(collection):
    class_name = collection.rna_type.__class__.__name__
    clean_name = class_name.replace("BlendData", "").lower()
    return clean_name

class RealTimeEngine(bpy.types.RenderEngine):
    bl_idname = 'RTE_FRAMEWORK'
    bl_label = "Real Time Engine Framework"

    def __init__(self, watch_list=DEFAULT_WATCHLIST):
        self._watch_list = [getattr(bpy.data, i) for i in watch_list]

        self._tracking_sets = {}
        for collection in self._watch_list:
            collection_name = get_collection_name(collection)
            self._tracking_sets[collection_name] = set()

        # Display image
        self.width = 1
        self.height = 1
        self.display = Buffer(GL_BYTE, 3, [0, 0, 0])

        self._old_vmat = None
        self._old_pmat = None
        self._old_viewport = None

        def main_loop(scene):
            try:
                self.scene_callback()
            except ReferenceError:
                bpy.apps.handlers.scene_update_post.remove(main_loop)

        bpy.app.handlers.scene_update_post.append(main_loop)

    def view_update(self, context):
        """ Called when the scene is changed """
        for collection in self._watch_list:
            collection_name = get_collection_name(collection)
            collection_set = set(collection)
            tracking_set = self._tracking_sets[collection_name]

            # Check for new items
            add_method = getattr(self, "add_"+collection_name, self.add_data)
            add_set = collection_set - tracking_set
            add_method(add_set)
            tracking_set |= add_set

            # Check for removed items
            remove_method = getattr(self, "remove_"+collection_name, self.remove_data)
            remove_set = tracking_set - collection_set
            remove_method(remove_set)
            tracking_set -= remove_set

            # Check for updates
            update_method = getattr(self, "update_"+collection_name, self.update_data)
            update_set = [item for item in collection if item.is_updated]
            update_method(update_set)

    def view_draw(self, context):
        """ Called when viewport settings change """
        region = context.region
        view = context.region_data

        vmat = view.view_matrix.copy()
        vmat_inv = vmat.inverted()
        pmat = view.perspective_matrix * vmat_inv

        viewport = [region.x, region.y, region.width, region.height]

        self.update_view(vmat, pmat, viewport)

        glGetError()
        glPushAttrib(GL_ALL_ATTRIB_BITS)

        glDisable(GL_DEPTH_TEST)
        glDisable(GL_CULL_FACE)
        glEnable(GL_TEXTURE_2D)

        glClearColor(0, 0, 1, 1)
        glClear(GL_COLOR_BUFFER_BIT)

        glMatrixMode(GL_MODELVIEW)
        glPushMatrix()
        glLoadIdentity()
        glMatrixMode(GL_PROJECTION)
        glPushMatrix()
        glLoadIdentity()

        bgl_buffer = Buffer(GL_INT, 1)
        glActiveTexture(GL_TEXTURE0)
        glGenTextures(1, bgl_buffer)
        tex = bgl_buffer[0]
        glBindTexture(GL_TEXTURE_2D, tex)
        width = viewport[2] - viewport[0]
        height = viewport[3] - viewport[1]
        glTexImage2D(GL_TEXTURE_2D, 0, 3, self.width, self.height, 0, GL_RGB, GL_BYTE, self.display)

        glBegin(GL_QUADS)
        glTexCoord2f(0.0, 0.0)
        glVertex3i(-1, -1, 0)
        glTexCoord2f(1.0, 0.0)
        glVertex3i(1, -1, 0)
        glTexCoord2f(1.0, 1.0)
        glVertex3i(1, 1, 0)
        glTexCoord2f(0.0, 1.0)
        glVertex3i(-1, 1, 0)
        glEnd()

        glDeleteTextures(1, bgl_buffer)

        glPopMatrix()
        glMatrixMode(GL_MODELVIEW)
        glPopMatrix()

        glPopAttrib()

    def add_objects(self, add_set):
        for data in add_set:
            print(data.__class__.__name__, data.name, "added with Object add")

    def remove_objects(self, add_set):
        for data in add_set:
            print(data.__class__.__name__, data.name, "removed with Object remove")

    def update_objects(self, add_set):
        for data in add_set:
            print(data.__class__.__name__, data.name, "updated with Object update")

    def add_data(self, add_set):
        for data in add_set:
            print(data.__class__.__name__, data.name, "added with generic add")

    def remove_data(self, remove_set):
        for data in remove_set:
            print(data.__class__.__name__, data.name, "removed with generic remove")

    def update_data(self, update_set):
        for data in update_set:
            print(data.__class__.__name__, data.name, "updated with generic update")

    def update_view(self, view_matrix, projection_matrix, viewport):
        print("update_view")
        if view_matrix != self._old_vmat:
            print(view_matrix)
        if projection_matrix != self._old_pmat:
            print(projection_matrix)
        if viewport != viewport:
            print(viewport)

    def scene_callback(self):
        pass

<|endoftext|>