<commit_before>import os, re, datetime, json, csv
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
        '''
        header_fixes = False
        if (kwargs["fasta_header_fix"]):
            header_fixes = {}
            try:
                with open(kwargs["fasta_header_fix"], 'rU') as fh:
                    for line in fh:
                        if not line.startswith('#'):
                            k, v = line.strip().split("\t")
                            header_fixes[k] = v
            except IOError:
                raise Exception(kwargs["fasta_header_fix"], "not found")

        viruses = []
        sequences = []
        try:
            handle = open(fasta, 'r')
        except IOError:
            raise Exception(fasta, "not found")
        else:
            for record in SeqIO.parse(handle, "fasta"):
                if header_fixes:
                    try:
                        record.description = header_fixes[record.description]
                    except KeyError:
                        raise Exception(record.description, "not in header fix file. Fatal.")
                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, title, 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 'title' not in v and title is not None:
            if title == 'null' or title == 'none' or title == 'None':
                v['title'] = None
            else:
                v['title'] = title
        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, n_entrez=2500, **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), n_entrez)): # split accessions list into 2500-long portions
            queries.append(" ".join(accessions[i:i+n_entrez])) # 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()
            # set up as none and overwrite
            s["title"] = None
            s['authors'] = None
            s["puburl"] = None
            s["journal"] = None
            print("Processing genbank file for " + s['accession'])
            # all the references (i.e. the papers / direct-submission notes)
            references = record.annotations["references"]

            if len(references):
                # is there a reference which is not a "Direct Submission"?
                titles = [reference.title for reference in references]
                try:
                    idx = [i for i, j in enumerate(titles) if j is not None and j != "Direct Submission"][0]
                except IndexError: # fall back to direct submission
                    idx = [i for i, j in enumerate(titles) if j is not None][0]
                reference = references[idx] # <class 'Bio.SeqFeature.Reference'>
                keys = reference.__dict__.keys()
                s['title'] = reference.title
                if "authors" in keys and reference.authors is not None:
                    first_author = re.match(r'^([^,]*)', reference.authors).group(0)
                    s['authors'] = first_author + " et al"
                if "journal" in keys and reference.journal is not None:
                    s['journal'] = reference.journal
                if "pubmed_id" in keys and reference.pubmed_id is not None:
                    s["puburl"] = "https://www.ncbi.nlm.nih.gov/pubmed/" + reference.pubmed_id
            else:
                print("Couldn't find the reference for " + s['accession'])

            # print(" *** Accession: {} title: {} authors: {} journal: {} paperURL: {}".format(s['accession'], s['title'], s['authors'], s['journal'], s['puburl']))

            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(str(len(viruses)) + " genbank entries parsed")
        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:
            if collection_date.split('-')[0].isdigit():
                if int(collection_date.split('-')[0]) < 13:
                    return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%m-%Y'), '%Y-%m-XX')
                else:
                    return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%Y-%m'), '%Y-%m-XX')
            else:
                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>Updating fauna parsing to only include puburl when present<commit_after>import os, re, datetime, json, csv
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
        '''
        header_fixes = False
        if (kwargs["fasta_header_fix"]):
            header_fixes = {}
            try:
                with open(kwargs["fasta_header_fix"], 'rU') as fh:
                    for line in fh:
                        if not line.startswith('#'):
                            k, v = line.strip().split("\t")
                            header_fixes[k] = v
            except IOError:
                raise Exception(kwargs["fasta_header_fix"], "not found")

        viruses = []
        sequences = []
        try:
            handle = open(fasta, 'r')
        except IOError:
            raise Exception(fasta, "not found")
        else:
            for record in SeqIO.parse(handle, "fasta"):
                if header_fixes:
                    try:
                        record.description = header_fixes[record.description]
                    except KeyError:
                        raise Exception(record.description, "not in header fix file. Fatal.")
                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, title, 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 'title' not in v and title is not None:
            if title == 'null' or title == 'none' or title == 'None':
                v['title'] = None
            else:
                v['title'] = title
        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, n_entrez=2500, **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), n_entrez)): # split accessions list into 2500-long portions
            queries.append(" ".join(accessions[i:i+n_entrez])) # 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()
            # set up as none and overwrite
            s["title"] = None
            s['authors'] = None
            s["puburl"] = None
            s["journal"] = None
            print("Processing genbank file for " + s['accession'])
            # all the references (i.e. the papers / direct-submission notes)
            references = record.annotations["references"]

            if len(references):
                # is there a reference which is not a "Direct Submission"?
                titles = [reference.title for reference in references]
                try:
                    idx = [i for i, j in enumerate(titles) if j is not None and j != "Direct Submission"][0]
                except IndexError: # fall back to direct submission
                    idx = [i for i, j in enumerate(titles) if j is not None][0]
                reference = references[idx] # <class 'Bio.SeqFeature.Reference'>
                keys = reference.__dict__.keys()
                s['title'] = reference.title
                if "authors" in keys and reference.authors:
                    first_author = re.match(r'^([^,]*)', reference.authors).group(0)
                    s['authors'] = first_author + " et al"
                if "journal" in keys and reference.journal:
                    s['journal'] = reference.journal
                if "pubmed_id" in keys and reference.pubmed_id:
                    s["puburl"] = "https://www.ncbi.nlm.nih.gov/pubmed/" + reference.pubmed_id
            else:
                print("Couldn't find the reference for " + s['accession'])

            # print(" *** Accession: {} title: {} authors: {} journal: {} paperURL: {}".format(s['accession'], s['title'], s['authors'], s['journal'], s['puburl']))

            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(str(len(viruses)) + " genbank entries parsed")
        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:
            if collection_date.split('-')[0].isdigit():
                if int(collection_date.split('-')[0]) < 13:
                    return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%m-%Y'), '%Y-%m-XX')
                else:
                    return datetime.datetime.strftime(datetime.datetime.strptime(collection_date,'%Y-%m'), '%Y-%m-XX')
            else:
                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># Copyright 2020 Camptocamp SA
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html)

from odoo import models

from odoo.addons.queue_job.job import job

class StockPicking(models.Model):
    _inherit = "stock.picking"

    def action_done(self):
        res = super().action_done()
        picking_to_invoice = self.filtered(
            lambda r: r.sale_id.partner_invoice_id.invoicing_mode == "at_shipping"
            and r.picking_type_code == "outgoing"
        )
        for picking in picking_to_invoice:
            picking.with_delay()._invoicing_at_shipping()
        return res

    @job(default_channel="root.invoice_at_shipping")
    def _invoicing_at_shipping(self):
        self.ensure_one()
        sale_order_ids = self._get_sales_order_to_invoice()
        partner = sale_order_ids.partner_invoice_id
        invoices = sale_order_ids._create_invoices(
            grouped=partner.one_invoice_per_order
        )
        for invoice in invoices:
            invoice.with_delay()._validate_invoice()
        return invoices

    def _get_sales_order_to_invoice(self):
        return self.mapped("move_lines.group_id.sale_id")
<commit_msg>Invoice sales based on moves<commit_after># Copyright 2020 Camptocamp SA
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html)

from odoo import models

from odoo.addons.queue_job.job import job

class StockPicking(models.Model):
    _inherit = "stock.picking"

    def action_done(self):
        res = super().action_done()
        picking_to_invoice = self.filtered(
            lambda r: r.sale_id.partner_invoice_id.invoicing_mode == "at_shipping"
            and r.picking_type_code == "outgoing"
        )
        for picking in picking_to_invoice:
            picking.with_delay()._invoicing_at_shipping()
        return res

    @job(default_channel="root.invoice_at_shipping")
    def _invoicing_at_shipping(self):
        self.ensure_one()
        sale_order_ids = self._get_sales_order_to_invoice()
        partner = sale_order_ids.partner_invoice_id
        invoices = sale_order_ids._create_invoices(
            grouped=partner.one_invoice_per_order
        )
        for invoice in invoices:
            invoice.with_delay()._validate_invoice()
        return invoices

    def _get_sales_order_to_invoice(self):
        return self.mapped("move_lines.sale_line_id.order_id")
<|endoftext|>


<commit_before># Amara, universalsubtitles.org
#
# Copyright (C) 2012 Participatory Culture Foundation
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see
# http://www.gnu.org/licenses/agpl-3.0.html.

from httplib2 import Http
from urllib import urlencode

from django.utils.translation import ugettext_lazy as _

from utils import send_templated_email
from utils.metrics import Meter
from libs.unilangs import LanguageCode
from videos.models import Video

import logging
logger = logging.getLogger("team-notifier")

class BaseNotification(object):
    """
    Holds the data needed to prepare a notification.
    Subclasses should be able to translate video_ids
    and language codes with from_internal_lang and
    from_internal_video_id

    Also, subclasses should implement a more specialized version of
    'send_http_request'
    'send_email'
    """
    codec = "unisubs"
    api_name = "partners"
    def from_internal_lang(self, lang_code):
        # we allow empty language codes
        if not lang_code:
            return ""
        return LanguageCode(lang_code, "unisubs").encode(self.codec)

    def to_internal_lang(self, lang_code):
        """
        """
        if not lang_code:
            return ""
        return LanguageCode(lang_code, self.codec).encode("unisubs")

    def to_internal_video_id(self, api_pk):
        """
        Coverts the public api id to the internal video_id for
        this api resource.

        Subclasses mapping to other system's id e.g ted should be able
        to gather the video from this public id
        """
        return Video.objects.get(video_id=api_pk).video_id

    def from_internal_video_id(self, video_id, video=None):
        """
        Coverts the internal video id representation (the actual )
        Video.video_id into a public video id. Partners can override
        the logic to fetch to their ids here.
        If a video has already been fetched from the db, it can be passed
        to avoid an extra lookup.
        """
        return video_id if video_id else video.video_id

    def __init__(self, team, partner, event_name,  **kwargs):
        """
        If the event is about new / edits to videos, then language_pk
        will be None else it can be about languages or subtitles.
        """
        from videos.models import Video
        self.team = team
        self.partner = partner
        video_id  = kwargs.pop('video_id', None)
        if video_id:
            self.video = Video.objects.get(video_id=video_id)
        else:
            self.video = None
        self.language_pk = kwargs.pop('language_pk', None)
        self.application_pk = kwargs.pop('application_pk', None)
        self.version_pk = kwargs.pop('version_pk', None)
        if self.language_pk:
            self.language = self.video.subtitlelanguage_set.get(pk=self.language_pk)
            if self.version_pk:
                self.version = self.language.subtitleversion_set.get(pk=self.version_pk)
        else:
            self.language = None
        self.event_name = event_name
        self.api_url = self.get_api_url()

    def get_api_url(self):
        """
        Returns what api url the recipient of this notification should
        query for the latest data. This is team dependent if the team
        has a custom base url.
        """
        from apiv2.api import (
            VideoLanguageResource, VideoResource, ApplicationResource
        )
        from teams.models import Application
        if self.video:
            video_klass = getattr(self.__class__, "video_resource_class", VideoResource)
            if self.language:
                lang_klass = getattr(self.__class__, "language_resource_class", VideoLanguageResource)
                url =  lang_klass(self.api_name).get_resource_uri(self.language)
                if self.version_pk:
                   url += "subtitles/?version_no=%s"  % self.version.version_no
                return url
            else:
                return video_klass(self.api_name).get_resource_uri(self.video)
        elif self.application_pk:
            return ApplicationResource("partners").get_resource_uri(Application.objects.get(pk=self.application_pk))

    @property
    def video_id(self):
        if self.video:
            return self.from_internal_video_id(None, video=self.video)

    @property
    def language_code(self):
        if self.language:
            return  self.from_internal_lang(self.language.language)

    def send_http_request(self, url, basic_auth_username, basic_auth_password):
        h = Http()
        if basic_auth_username and basic_auth_password:
            h.add_credentials(basic_auth_username, basic_auth_password)

        project = self.video.get_team_video().project.slug if self.video else None
        data = {
            'event': self.event_name,
            'api_url': self.api_url,
        }
        if self.team:
            data['team'] = self.team.slug
        if self.partner:
            data['partner'] = self.partner.slug
        if project:
            data['project'] = project
        if self.video:
            data['video_id'] = self.video_id
        if self.application_pk:
            data['application_id'] = self.application_pk
        if self.language_code:
            data.update({"language_code":self.language_code} )
        data_sent = data
        data = urlencode(data)
        url = "%s?%s"  % (url , data)
        try:
            resp, content = h.request(url, method="POST", body=data)
            success =  200<= resp.status <400
            if success is False:
                logger.error("Failed to notify team %s " %( self.team),
                     extra={
                        'team': self.team or self.partner,
                        'url': url,
                        'response': resp,
                        'content': content,
                        'data_sent':data_sent,
                    })

                Meter('http-callback-notification-error').inc()
            else:
                Meter('http-callback-notification-success').inc()
            return success, content
        except:
            logger.exception("Failed to send http notification ")
        return None, None

    def send_email(self, email_to):
        Meter('templated-emails-sent-by-type.teams.team-video-activity').inc()
        send_templated_email(email_to,
                _("New activity on your team video"),
                "teams/emails/new-activity.html",
                {
                    "video": self.video,
                    "event-name": self.event_name,
                    "team": self.team,
                    "laguage":self.language,
                }
            )
<commit_msg>Fix bad whitespace.<commit_after># Amara, universalsubtitles.org
#
# Copyright (C) 2012 Participatory Culture Foundation
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see
# http://www.gnu.org/licenses/agpl-3.0.html.

from httplib2 import Http
from urllib import urlencode

from django.utils.translation import ugettext_lazy as _

from utils import send_templated_email
from utils.metrics import Meter
from libs.unilangs import LanguageCode
from videos.models import Video

import logging
logger = logging.getLogger("team-notifier")

class BaseNotification(object):
    """
    Holds the data needed to prepare a notification.
    Subclasses should be able to translate video_ids
    and language codes with from_internal_lang and
    from_internal_video_id

    Also, subclasses should implement a more specialized version of
    'send_http_request'
    'send_email'
    """
    codec = "unisubs"
    api_name = "partners"
    def from_internal_lang(self, lang_code):
        # we allow empty language codes
        if not lang_code:
            return ""
        return LanguageCode(lang_code, "unisubs").encode(self.codec)

    def to_internal_lang(self, lang_code):
        """
        """
        if not lang_code:
            return ""
        return LanguageCode(lang_code, self.codec).encode("unisubs")

    def to_internal_video_id(self, api_pk):
        """
        Coverts the public api id to the internal video_id for
        this api resource.

        Subclasses mapping to other system's id e.g ted should be able
        to gather the video from this public id
        """
        return Video.objects.get(video_id=api_pk).video_id

    def from_internal_video_id(self, video_id, video=None):
        """
        Coverts the internal video id representation (the actual )
        Video.video_id into a public video id. Partners can override
        the logic to fetch to their ids here.
        If a video has already been fetched from the db, it can be passed
        to avoid an extra lookup.
        """
        return video_id if video_id else video.video_id

    def __init__(self, team, partner, event_name,  **kwargs):
        """
        If the event is about new / edits to videos, then language_pk
        will be None else it can be about languages or subtitles.
        """
        from videos.models import Video
        self.team = team
        self.partner = partner
        video_id  = kwargs.pop('video_id', None)
        if video_id:
            self.video = Video.objects.get(video_id=video_id)
        else:
            self.video = None
        self.language_pk = kwargs.pop('language_pk', None)
        self.application_pk = kwargs.pop('application_pk', None)
        self.version_pk = kwargs.pop('version_pk', None)
        if self.language_pk:
            self.language = self.video.subtitlelanguage_set.get(pk=self.language_pk)
            if self.version_pk:
                self.version = self.language.subtitleversion_set.get(pk=self.version_pk)
        else:
            self.language = None
        self.event_name = event_name
        self.api_url = self.get_api_url()

    def get_api_url(self):
        """
        Returns what api url the recipient of this notification should
        query for the latest data. This is team dependent if the team
        has a custom base url.
        """
        from apiv2.api import (
            VideoLanguageResource, VideoResource, ApplicationResource
        )
        from teams.models import Application
        if self.video:
            video_klass = getattr(self.__class__, "video_resource_class", VideoResource)
            if self.language:
                lang_klass = getattr(self.__class__, "language_resource_class", VideoLanguageResource)
                url =  lang_klass(self.api_name).get_resource_uri(self.language)
                if self.version_pk:
                   url += "subtitles/?version_no=%s"  % self.version.version_no
                return url
            else:
                return video_klass(self.api_name).get_resource_uri(self.video)
        elif self.application_pk:
            return ApplicationResource("partners").get_resource_uri(Application.objects.get(pk=self.application_pk))

    @property
    def video_id(self):
        if self.video:
            return self.from_internal_video_id(None, video=self.video)

    @property
    def language_code(self):
        if self.language:
            return  self.from_internal_lang(self.language.language)

    def send_http_request(self, url, basic_auth_username, basic_auth_password):
        h = Http()
        if basic_auth_username and basic_auth_password:
            h.add_credentials(basic_auth_username, basic_auth_password)

        project = self.video.get_team_video().project.slug if self.video else None
        data = {
            'event': self.event_name,
            'api_url': self.api_url,
        }
        if self.team:
            data['team'] = self.team.slug
        if self.partner:
            data['partner'] = self.partner.slug
        if project:
            data['project'] = project
        if self.video:
            data['video_id'] = self.video_id
        if self.application_pk:
            data['application_id'] = self.application_pk
        if self.language_code:
            data.update({"language_code":self.language_code} )
        data_sent = data
        data = urlencode(data)
        url = "%s?%s" % (url , data)
        try:
            resp, content = h.request(url, method="POST", body=data)
            success = 200 <= resp.status < 400
            if success is False:
                logger.error("Failed to notify team %s " % (self.team),
                     extra={
                        'team': self.team or self.partner,
                        'url': url,
                        'response': resp,
                        'content': content,
                        'data_sent':data_sent,
                    })

                Meter('http-callback-notification-error').inc()
            else:
                Meter('http-callback-notification-success').inc()
            return success, content
        except:
            logger.exception("Failed to send http notification ")
        return None, None

    def send_email(self, email_to):
        Meter('templated-emails-sent-by-type.teams.team-video-activity').inc()
        send_templated_email(email_to,
                _("New activity on your team video"),
                "teams/emails/new-activity.html",
                {
                    "video": self.video,
                    "event-name": self.event_name,
                    "team": self.team,
                    "laguage":self.language,
                }
            )

<|endoftext|>


<commit_before>#!/opt/local/bin/python2.5
"""
An implementation of:
Improving Translation Quality by Discarding Most of the Phrasetable (Johnson, Martin, Foster, Kuhn 2007)

Usage:
    -d debug symbols
    -i input file
    -o ouput file
"""
# Copyright (C) 2009
# Authors: Gabriel Synnaeve (& Nicolas Dumazet)
# License: http://www.opensource.org/licenses/PythonSoftFoundation.php

import sys, getopt, math #, guppy
#h = guppy.hpy()

import enrichment

def Usage():
    print "./pruning.py phrase-table [-d][-h][-o outputfile]"
    sys.exit(0)

try:
    opts, args = getopt.getopt(sys.argv[1:],'di:o:h')
except getopt.GetoptError:
    Usage()
    sys.exit(2)
debug = 0
for o, a in opts:
    if o == '-h':
        Usage()
    elif o == '-d':
        debug = 1
    elif o == '-i':
        inputname = a
    elif o == '-o':
        outputname = a
    else:
        assert False, "unhandled option"

#file = open(sys.argv[1])
file = open(inputname)

count_s = {}
count_t = {}
dict_st = {}
lines = {}

def count(line):
    table = line.replace('#','').replace('[','').replace(']','')\
            .replace('?','').replace('!','').replace(',','')\
            .replace("'",'').replace('&quo','').replace(';','')\
            .replace('--','')\
            .replace('  ',' ').strip().split('|||')
            # seems to be faster than with the RE '#\[\]?!'
    source = table[0].strip()
    target = table[1].strip()
    if not source == '' and not target == '':
        if source in count_s:
            count_s[source] +=  1
        else:
            count_s[source] =  1
        if target in count_t:
            count_t[target] +=  1
        else:
            count_t[target] =  1
        if source in dict_st:
            d = dict_st[source]
            if target in d:
                d[target] += 1
            else:
                d[target] = 1
        else:
            dict_st[source] = {target:1}
        if (source, target) in lines:
            lines[(source, target)].append(N)
        else:
            lines[(source, target)] = [N]

#map(count, file)
N = 0
for line in file:
    N += 1
    count(line)

included_in_s = {}
included_in_t = {}
sets = {}

for k in count_s.iterkeys():
    tmp = 0
    dico = {}
    for (s, v) in count_s.iteritems():
        if k in s and k != s:
            tmp += v
            dico[s] = True
            if not s in sets:
                sets[s] = set(dict_st[s].keys())
    included_in_s[k] = dico
    count_s[k] += tmp

for k in count_t.iterkeys():
    tmp = 0
    dico = {}
    for (t, v) in count_t.iteritems():
        if k in t and k != t:
            tmp += v
            dico[t] = True
    included_in_t[k] = dico
    count_t[k] += tmp

for ks, kdic in dict_st.iteritems():
    ks_lesser_than = included_in_s[ks]
    for kt in kdic.iterkeys():
        # kt is in every element of kt_lesser_than
        kt_lesser_than = included_in_t[kt]
        kt_lesser_than = set(kt_lesser_than.keys())
        for s in ks_lesser_than.iterkeys():
            # we know that 'ks in s'
            intersect = kt_lesser_than & sets[s]
            tmp = 0
            for e in intersect:
                tmp += dict_st[s][e]
            kdic[kt] += tmp

file.close()

if debug:
    #print h.heap()
    print '===========count_s============='
    print count_s
    print '===========count_t============='
    print count_t
    print '===========dict_st============='
    print dict_st
    print '============================'

delete = []
threshold = math.log(N) - 0.01
if debug:
    print ">>> Threshold :",
    print threshold
##################################
for ks, kdic in dict_st.iteritems():
    for kt in kdic.iterkeys():
        try:
            #print -math.log( fisher_exact_text(dict_st[ks][kt], count_s[ks], \
            #        count_t[kt], N) [1] )
            #print k
            print "dict_st[ks][kt] ",
            print dict_st[ks][kt]
            print "count_s[ks] ",
            print count_s[ks]
            print 'debut|'+ks+'|fin'
            print "count_t[kt] ",
            print count_t[kt]
            print 'debut|'+kt+'|fin'
            print "N ",
            print N
            if -math.log( enrichment.fisher_exact_test(dict_st[ks][kt], \
                    count_s[ks], count_t[kt], N)[1] ) > threshold:
                for l in lines[(ks, kt)]:
                    delete.append(l)
        except OverflowError:
            # The value is so low that even your mo****
            # Discard this entry
            for l in lines[(ks, kt)]:
                delete.append(l)

delete.sort()

if debug:
    print ">>> Lines to delete:",
    print delete

file = open(inputname)
outputfile = open(outputname, 'w')
i = 0
for line in file:
    i += 1
    if not i in delete:
        outputfile.write(line)
outputfile.close()

<commit_msg>some trick with whitespaces for substring matching<commit_after>#!/opt/local/bin/python2.5
"""
An implementation of:
Improving Translation Quality by Discarding Most of the Phrasetable (Johnson, Martin, Foster, Kuhn 2007)

Usage:
    -d debug symbols
    -i input file
    -o ouput file
"""
# Copyright (C) 2009
# Authors: Gabriel Synnaeve (& Nicolas Dumazet)
# License: http://www.opensource.org/licenses/PythonSoftFoundation.php

import sys, getopt, math #, guppy
#h = guppy.hpy()

import enrichment

def Usage():
    print "./pruning.py phrase-table [-d][-h][-o outputfile]"
    sys.exit(0)

try:
    opts, args = getopt.getopt(sys.argv[1:],'di:o:h')
except getopt.GetoptError:
    Usage()
    sys.exit(2)
debug = 0
for o, a in opts:
    if o == '-h':
        Usage()
    elif o == '-d':
        debug = 1
    elif o == '-i':
        inputname = a
    elif o == '-o':
        outputname = a
    else:
        assert False, "unhandled option"

#file = open(sys.argv[1])
file = open(inputname)

count_s = {}
count_t = {}
dict_st = {}
lines = {}

def count(line):
    table = line.replace('#','').replace('[','').replace(']','')\
            .replace('?','').replace('!','').replace(',','')\
            .replace("'",'').replace('&quo','').replace(';','')\
            .replace('--','')\
            .replace('  ',' ').strip().split('|||')
            # seems to be faster than with the RE '#\[\]?!'
    table[0] = ' '+table[0]+' '
    table[1] = ' '+table[1]+' '
    source = table[0]
    target = table[1]
    if not source == '' and not target == '':
        if source in count_s:
            count_s[source] +=  1
        else:
            count_s[source] =  1
        if target in count_t:
            count_t[target] +=  1
        else:
            count_t[target] =  1
        if source in dict_st:
            d = dict_st[source]
            if target in d:
                d[target] += 1
            else:
                d[target] = 1
        else:
            dict_st[source] = {target:1}
        if (source, target) in lines:
            lines[(source, target)].append(N)
        else:
            lines[(source, target)] = [N]

#map(count, file)
N = 0
for line in file:
    N += 1
    count(line)

included_in_s = {}
included_in_t = {}
sets = {}

for k in count_s.iterkeys():
    tmp = 0
    dico = {}
    for (s, v) in count_s.iteritems():
        if ' '+k+' ' in s and k != s:
            tmp += v
            dico[s] = True
            if not s in sets:
                sets[s] = set(dict_st[s].keys())
    included_in_s[k] = dico
    count_s[k] += tmp

for k in count_t.iterkeys():
    tmp = 0
    dico = {}
    for (t, v) in count_t.iteritems():
        if ' '+k+' ' in t and k != t:
            tmp += v
            dico[t] = True
    included_in_t[k] = dico
    count_t[k] += tmp

for ks, kdic in dict_st.iteritems():
    ks_lesser_than = included_in_s[ks]
    for kt in kdic.iterkeys():
        # kt is in every element of kt_lesser_than
        kt_lesser_than = included_in_t[kt]
        kt_lesser_than = set(kt_lesser_than.keys())
        for s in ks_lesser_than.iterkeys():
            # we know that 'ks in s'
            intersect = kt_lesser_than & sets[s]
            tmp = 0
            for e in intersect:
                tmp += dict_st[s][e]
            kdic[kt] += tmp

file.close()

if debug:
    #print h.heap()
    print '===========count_s============='
    print count_s
    print '===========count_t============='
    print count_t
    print '===========dict_st============='
    print dict_st
    print '============================'

delete = []
threshold = math.log(N) - 0.01
if debug:
    print ">>> Threshold :",
    print threshold
##################################
for ks, kdic in dict_st.iteritems():
    for kt in kdic.iterkeys():
        try:
            if debug:
                print "dict_st[ks][kt] ",
                print dict_st[ks][kt]
                print "count_s[ks] ",
                print count_s[ks]
                print 'debut|'+ks+'|fin'
                print "count_t[kt] ",
                print count_t[kt]
                print 'debut|'+kt+'|fin'
                print "N ",
                print N
            if -math.log( enrichment.fisher_exact_test(dict_st[ks][kt], \
                    count_s[ks], count_t[kt], N)[1] ) > threshold:
                for l in lines[(ks, kt)]:
                    delete.append(l)
        except OverflowError:
            # The value is so low that even your mo****
            # Discard this entry
            for l in lines[(ks, kt)]:
                delete.append(l)

delete.sort()

if debug:
    print ">>> Lines to delete:",
    print delete

file = open(inputname)
outputfile = open(outputname, 'w')
i = 0
for line in file:
    i += 1
    if not i in delete:
        outputfile.write(line)
outputfile.close()

<|endoftext|>


<commit_before># Copyright 2022 Katherine Crowson and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import warnings
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import torch

from scipy import integrate

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, deprecate
from .scheduling_utils import SchedulerMixin

@dataclass
class LMSDiscreteSchedulerOutput(BaseOutput):
    """
    Output class for the scheduler's step function output.

    Args:
        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            The predicted denoised sample (x_{0}) based on the model output from the current timestep.
            `pred_original_sample` can be used to preview progress or for guidance.
    """

    prev_sample: torch.FloatTensor
    pred_original_sample: Optional[torch.FloatTensor] = None

class LMSDiscreteScheduler(SchedulerMixin, ConfigMixin):
    """
    Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by
    Katherine Crowson:
    https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181

    [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__`
    function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`.
    [`~ConfigMixin`] also provides general loading and saving functionality via the [`~ConfigMixin.save_config`] and
    [`~ConfigMixin.from_config`] functions.

    Args:
        num_train_timesteps (`int`): number of diffusion steps used to train the model.
        beta_start (`float`): the starting `beta` value of inference.
        beta_end (`float`): the final `beta` value.
        beta_schedule (`str`):
            the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
            `linear` or `scaled_linear`.
        trained_betas (`np.ndarray`, optional):
            option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc.

    """

    @register_to_config
    def __init__(
        self,
        num_train_timesteps: int = 1000,
        beta_start: float = 0.0001,
        beta_end: float = 0.02,
        beta_schedule: str = "linear",
        trained_betas: Optional[np.ndarray] = None,
        **kwargs,
    ):
        deprecate(
            "tensor_format",
            "0.5.0",
            "If you're running your code in PyTorch, you can safely remove this argument.",
            take_from=kwargs,
        )

        if trained_betas is not None:
            self.betas = torch.from_numpy(trained_betas)
        elif beta_schedule == "linear":
            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
        elif beta_schedule == "scaled_linear":
            # this schedule is very specific to the latent diffusion model.
            self.betas = (
                torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
            )
        else:
            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")

        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)

        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
        sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32)
        self.sigmas = torch.from_numpy(sigmas)

        # standard deviation of the initial noise distribution
        self.init_noise_sigma = self.sigmas.max()

        # setable values
        self.num_inference_steps = None
        timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy()
        self.timesteps = torch.from_numpy(timesteps)
        self.derivatives = []
        self.is_scale_input_called = False

    def scale_model_input(
        self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor]
    ) -> torch.FloatTensor:
        """
        Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the K-LMS algorithm.

        Args:
            sample (`torch.FloatTensor`): input sample
            timestep (`float` or `torch.FloatTensor`): the current timestep in the diffusion chain

        Returns:
            `torch.FloatTensor`: scaled input sample
        """
        if isinstance(timestep, torch.Tensor):
            timestep = timestep.to(self.timesteps.device)
        step_index = (self.timesteps == timestep).nonzero().item()
        sigma = self.sigmas[step_index]
        sample = sample / ((sigma**2 + 1) ** 0.5)
        self.is_scale_input_called = True
        return sample

    def get_lms_coefficient(self, order, t, current_order):
        """
        Compute a linear multistep coefficient.

        Args:
            order (TODO):
            t (TODO):
            current_order (TODO):
        """

        def lms_derivative(tau):
            prod = 1.0
            for k in range(order):
                if current_order == k:
                    continue
                prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k])
            return prod

        integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0]

        return integrated_coeff

    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
        """
        Sets the timesteps used for the diffusion chain. Supporting function to be run before inference.

        Args:
            num_inference_steps (`int`):
                the number of diffusion steps used when generating samples with a pre-trained model.
            device (`str` or `torch.device`, optional):
                the device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
        """
        self.num_inference_steps = num_inference_steps

        timesteps = np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps, dtype=float)[::-1].copy()
        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
        sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
        sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32)
        self.sigmas = torch.from_numpy(sigmas).to(device=device)
        self.timesteps = torch.from_numpy(timesteps).to(device=device)

        self.derivatives = []

    def step(
        self,
        model_output: torch.FloatTensor,
        timestep: Union[float, torch.FloatTensor],
        sample: torch.FloatTensor,
        order: int = 4,
        return_dict: bool = True,
    ) -> Union[LMSDiscreteSchedulerOutput, Tuple]:
        """
        Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
        process from the learned model outputs (most often the predicted noise).

        Args:
            model_output (`torch.FloatTensor`): direct output from learned diffusion model.
            timestep (`float`): current timestep in the diffusion chain.
            sample (`torch.FloatTensor`):
                current instance of sample being created by diffusion process.
            order: coefficient for multi-step inference.
            return_dict (`bool`): option for returning tuple rather than LMSDiscreteSchedulerOutput class

        Returns:
            [`~schedulers.scheduling_utils.LMSDiscreteSchedulerOutput`] or `tuple`:
            [`~schedulers.scheduling_utils.LMSDiscreteSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`.
            When returning a tuple, the first element is the sample tensor.

        """
        if not self.is_scale_input_called:
            warnings.warn(
                "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
                "See `StableDiffusionPipeline` for a usage example."
            )

        if isinstance(timestep, torch.Tensor):
            timestep = timestep.to(self.timesteps.device)
        if (
            isinstance(timestep, int)
            or isinstance(timestep, torch.IntTensor)
            or isinstance(timestep, torch.LongTensor)
        ):
            warnings.warn(
                "Integer timesteps in `LMSDiscreteScheduler.step()` are deprecated and will be removed in version"
                " 0.5.0. Make sure to pass one of the `scheduler.timesteps`."
            )
            step_index = timestep
        else:
            step_index = (self.timesteps == timestep).nonzero().item()
        sigma = self.sigmas[step_index]

        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
        pred_original_sample = sample - sigma * model_output

        # 2. Convert to an ODE derivative
        derivative = (sample - pred_original_sample) / sigma
        self.derivatives.append(derivative)
        if len(self.derivatives) > order:
            self.derivatives.pop(0)

        # 3. Compute linear multistep coefficients
        order = min(step_index + 1, order)
        lms_coeffs = [self.get_lms_coefficient(order, step_index, curr_order) for curr_order in range(order)]

        # 4. Compute previous sample based on the derivatives path
        prev_sample = sample + sum(
            coeff * derivative for coeff, derivative in zip(lms_coeffs, reversed(self.derivatives))
        )

        if not return_dict:
            return (prev_sample,)

        return LMSDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)

    def add_noise(
        self,
        original_samples: torch.FloatTensor,
        noise: torch.FloatTensor,
        timesteps: torch.FloatTensor,
    ) -> torch.FloatTensor:
        sigmas = self.sigmas.to(original_samples.device)
        schedule_timesteps = self.timesteps.to(original_samples.device)
        timesteps = timesteps.to(original_samples.device)
        if isinstance(timesteps, torch.IntTensor) or isinstance(timesteps, torch.LongTensor):
            warnings.warn(
                "Integer timesteps in `LMSDiscreteScheduler.add_noise()` are deprecated and will be removed in"
                " version 0.5.0. Make sure to pass values from `scheduler.timesteps`."
            )
            step_indices = timesteps
        else:
            step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]

        sigma = sigmas[step_indices].flatten()
        while len(sigma.shape) < len(original_samples.shape):
            sigma = sigma.unsqueeze(-1)

        noisy_samples = original_samples + noise * sigma
        return noisy_samples

    def __len__(self):
        return self.config.num_train_timesteps
<commit_msg>Better steps deprecation for LMS (#753)<commit_after># Copyright 2022 Katherine Crowson and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import warnings
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import torch

from scipy import integrate

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, deprecate
from .scheduling_utils import SchedulerMixin

@dataclass
class LMSDiscreteSchedulerOutput(BaseOutput):
    """
    Output class for the scheduler's step function output.

    Args:
        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            The predicted denoised sample (x_{0}) based on the model output from the current timestep.
            `pred_original_sample` can be used to preview progress or for guidance.
    """

    prev_sample: torch.FloatTensor
    pred_original_sample: Optional[torch.FloatTensor] = None

class LMSDiscreteScheduler(SchedulerMixin, ConfigMixin):
    """
    Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by
    Katherine Crowson:
    https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181

    [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__`
    function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`.
    [`~ConfigMixin`] also provides general loading and saving functionality via the [`~ConfigMixin.save_config`] and
    [`~ConfigMixin.from_config`] functions.

    Args:
        num_train_timesteps (`int`): number of diffusion steps used to train the model.
        beta_start (`float`): the starting `beta` value of inference.
        beta_end (`float`): the final `beta` value.
        beta_schedule (`str`):
            the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
            `linear` or `scaled_linear`.
        trained_betas (`np.ndarray`, optional):
            option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc.

    """

    @register_to_config
    def __init__(
        self,
        num_train_timesteps: int = 1000,
        beta_start: float = 0.0001,
        beta_end: float = 0.02,
        beta_schedule: str = "linear",
        trained_betas: Optional[np.ndarray] = None,
        **kwargs,
    ):
        deprecate(
            "tensor_format",
            "0.5.0",
            "If you're running your code in PyTorch, you can safely remove this argument.",
            take_from=kwargs,
        )

        if trained_betas is not None:
            self.betas = torch.from_numpy(trained_betas)
        elif beta_schedule == "linear":
            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
        elif beta_schedule == "scaled_linear":
            # this schedule is very specific to the latent diffusion model.
            self.betas = (
                torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
            )
        else:
            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")

        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)

        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
        sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32)
        self.sigmas = torch.from_numpy(sigmas)

        # standard deviation of the initial noise distribution
        self.init_noise_sigma = self.sigmas.max()

        # setable values
        self.num_inference_steps = None
        timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy()
        self.timesteps = torch.from_numpy(timesteps)
        self.derivatives = []
        self.is_scale_input_called = False

    def scale_model_input(
        self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor]
    ) -> torch.FloatTensor:
        """
        Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the K-LMS algorithm.

        Args:
            sample (`torch.FloatTensor`): input sample
            timestep (`float` or `torch.FloatTensor`): the current timestep in the diffusion chain

        Returns:
            `torch.FloatTensor`: scaled input sample
        """
        if isinstance(timestep, torch.Tensor):
            timestep = timestep.to(self.timesteps.device)
        step_index = (self.timesteps == timestep).nonzero().item()
        sigma = self.sigmas[step_index]
        sample = sample / ((sigma**2 + 1) ** 0.5)
        self.is_scale_input_called = True
        return sample

    def get_lms_coefficient(self, order, t, current_order):
        """
        Compute a linear multistep coefficient.

        Args:
            order (TODO):
            t (TODO):
            current_order (TODO):
        """

        def lms_derivative(tau):
            prod = 1.0
            for k in range(order):
                if current_order == k:
                    continue
                prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k])
            return prod

        integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0]

        return integrated_coeff

    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
        """
        Sets the timesteps used for the diffusion chain. Supporting function to be run before inference.

        Args:
            num_inference_steps (`int`):
                the number of diffusion steps used when generating samples with a pre-trained model.
            device (`str` or `torch.device`, optional):
                the device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
        """
        self.num_inference_steps = num_inference_steps

        timesteps = np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps, dtype=float)[::-1].copy()
        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
        sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
        sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32)
        self.sigmas = torch.from_numpy(sigmas).to(device=device)
        self.timesteps = torch.from_numpy(timesteps).to(device=device)

        self.derivatives = []

    def step(
        self,
        model_output: torch.FloatTensor,
        timestep: Union[float, torch.FloatTensor],
        sample: torch.FloatTensor,
        order: int = 4,
        return_dict: bool = True,
    ) -> Union[LMSDiscreteSchedulerOutput, Tuple]:
        """
        Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
        process from the learned model outputs (most often the predicted noise).

        Args:
            model_output (`torch.FloatTensor`): direct output from learned diffusion model.
            timestep (`float`): current timestep in the diffusion chain.
            sample (`torch.FloatTensor`):
                current instance of sample being created by diffusion process.
            order: coefficient for multi-step inference.
            return_dict (`bool`): option for returning tuple rather than LMSDiscreteSchedulerOutput class

        Returns:
            [`~schedulers.scheduling_utils.LMSDiscreteSchedulerOutput`] or `tuple`:
            [`~schedulers.scheduling_utils.LMSDiscreteSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`.
            When returning a tuple, the first element is the sample tensor.

        """
        if not self.is_scale_input_called:
            warnings.warn(
                "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
                "See `StableDiffusionPipeline` for a usage example."
            )

        if isinstance(timestep, torch.Tensor):
            timestep = timestep.to(self.timesteps.device)
        if (
            isinstance(timestep, int)
            or isinstance(timestep, torch.IntTensor)
            or isinstance(timestep, torch.LongTensor)
        ):
            deprecate(
                "timestep as an index",
                "0.5.0",
                "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
                " `LMSDiscreteScheduler.step()` will not be supported in future versions. Make sure to pass"
                " one of the `scheduler.timesteps` as a timestep.",
                standard_warn=False,
            )
            step_index = timestep
        else:
            step_index = (self.timesteps == timestep).nonzero().item()
        sigma = self.sigmas[step_index]

        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
        pred_original_sample = sample - sigma * model_output

        # 2. Convert to an ODE derivative
        derivative = (sample - pred_original_sample) / sigma
        self.derivatives.append(derivative)
        if len(self.derivatives) > order:
            self.derivatives.pop(0)

        # 3. Compute linear multistep coefficients
        order = min(step_index + 1, order)
        lms_coeffs = [self.get_lms_coefficient(order, step_index, curr_order) for curr_order in range(order)]

        # 4. Compute previous sample based on the derivatives path
        prev_sample = sample + sum(
            coeff * derivative for coeff, derivative in zip(lms_coeffs, reversed(self.derivatives))
        )

        if not return_dict:
            return (prev_sample,)

        return LMSDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)

    def add_noise(
        self,
        original_samples: torch.FloatTensor,
        noise: torch.FloatTensor,
        timesteps: torch.FloatTensor,
    ) -> torch.FloatTensor:
        sigmas = self.sigmas.to(original_samples.device)
        schedule_timesteps = self.timesteps.to(original_samples.device)
        timesteps = timesteps.to(original_samples.device)
        if isinstance(timesteps, torch.IntTensor) or isinstance(timesteps, torch.LongTensor):
            deprecate(
                "timesteps as indices",
                "0.5.0",
                "Passing integer indices  (e.g. from `enumerate(timesteps)`) as timesteps to"
                " `LMSDiscreteScheduler.add_noise()` will not be supported in future versions. Make sure to"
                " pass values from `scheduler.timesteps` as timesteps.",
                standard_warn=False,
            )
            step_indices = timesteps
        else:
            step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]

        sigma = sigmas[step_indices].flatten()
        while len(sigma.shape) < len(original_samples.shape):
            sigma = sigma.unsqueeze(-1)

        noisy_samples = original_samples + noise * sigma
        return noisy_samples

    def __len__(self):
        return self.config.num_train_timesteps
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2015, 2016 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import functools
import time
import gc

from twisted.internet import reactor

from .metric import (
    CounterMetric, CallbackMetric, DistributionMetric, CacheMetric,
    MemoryUsageMetric,
)

logger = logging.getLogger(__name__)

all_metrics = []
all_collectors = []

class Metrics(object):
    """ A single Metrics object gives a (mutable) slice view of the all_metrics
    dict, allowing callers to easily register new metrics that are namespaced
    nicely."""

    def __init__(self, name):
        self.name_prefix = name

    def register_collector(self, func):
        all_collectors.append(func)

    def _register(self, metric_class, name, *args, **kwargs):
        full_name = "%s_%s" % (self.name_prefix, name)

        metric = metric_class(full_name, *args, **kwargs)

        all_metrics.append(metric)
        return metric

    def register_counter(self, *args, **kwargs):
        return self._register(CounterMetric, *args, **kwargs)

    def register_callback(self, *args, **kwargs):
        return self._register(CallbackMetric, *args, **kwargs)

    def register_distribution(self, *args, **kwargs):
        return self._register(DistributionMetric, *args, **kwargs)

    def register_cache(self, *args, **kwargs):
        return self._register(CacheMetric, *args, **kwargs)

def register_memory_metrics(hs):
    try:
        import psutil
        process = psutil.Process()
        process.memory_info().rss
    except (ImportError, AttributeError):
        logger.warn(
            "psutil is not installed or incorrect version."
            " Disabling memory metrics."
        )
        return
    metric = MemoryUsageMetric(hs, psutil)
    all_metrics.append(metric)

def get_metrics_for(pkg_name):
    """ Returns a Metrics instance for conveniently creating metrics
    namespaced with the given name prefix. """

    # Convert a "package.name" to "package_name" because Prometheus doesn't
    # let us use . in metric names
    return Metrics(pkg_name.replace(".", "_"))

def render_all():
    strs = []

    for collector in all_collectors:
        collector()

    for metric in all_metrics:
        try:
            strs += metric.render()
        except Exception:
            strs += ["# FAILED to render"]
            logger.exception("Failed to render metric")

    strs.append("")  # to generate a final CRLF

    return "\n".join(strs)

reactor_metrics = get_metrics_for("reactor")
tick_time = reactor_metrics.register_distribution("tick_time")
pending_calls_metric = reactor_metrics.register_distribution("pending_calls")

gc_time = reactor_metrics.register_distribution("gc_time", labels=["gen"])
gc_unreachable = reactor_metrics.register_counter("gc_unreachable", labels=["gen"])

reactor_metrics.register_callback(
    "gc_counts", lambda: {(i,): v for i, v in enumerate(gc.get_count())}, labels=["gen"]
)

def runUntilCurrentTimer(func):

    @functools.wraps(func)
    def f(*args, **kwargs):
        now = reactor.seconds()
        num_pending = 0

        # _newTimedCalls is one long list of *all* pending calls. Below loop
        # is based off of impl of reactor.runUntilCurrent
        for delayed_call in reactor._newTimedCalls:
            if delayed_call.time > now:
                break

            if delayed_call.delayed_time > 0:
                continue

            num_pending += 1

        num_pending += len(reactor.threadCallQueue)

        start = time.time() * 1000
        ret = func(*args, **kwargs)
        end = time.time() * 1000
        tick_time.inc_by(end - start)
        pending_calls_metric.inc_by(num_pending)

        # Check if we need to do a manual GC (since its been disabled), and do
        # one if necessary.
        threshold = gc.get_threshold()
        counts = gc.get_count()
        for i in (2, 1, 0):
            if threshold[i] < counts[i]:
                logger.info("Collecting gc %d", i)

                start = time.time() * 1000
                unreachable = gc.collect(i)
                end = time.time() * 1000

                gc_time.inc_by(end - start, i)
                gc_unreachable.inc_by(unreachable, i)

        return ret

    return f

try:
    # Ensure the reactor has all the attributes we expect
    reactor.runUntilCurrent
    reactor._newTimedCalls
    reactor.threadCallQueue

    # runUntilCurrent is called when we have pending calls. It is called once
    # per iteratation after fd polling.
    reactor.runUntilCurrent = runUntilCurrentTimer(reactor.runUntilCurrent)

    # We manually run the GC each reactor tick so that we can get some metrics
    # about time spent doing GC,
    gc.disable()
except AttributeError:
    pass
<commit_msg>Allow creation of a 'subspace' within a Metrics object, returning another one<commit_after># -*- coding: utf-8 -*-
# Copyright 2015, 2016 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import functools
import time
import gc

from twisted.internet import reactor

from .metric import (
    CounterMetric, CallbackMetric, DistributionMetric, CacheMetric,
    MemoryUsageMetric,
)

logger = logging.getLogger(__name__)

all_metrics = []
all_collectors = []

class Metrics(object):
    """ A single Metrics object gives a (mutable) slice view of the all_metrics
    dict, allowing callers to easily register new metrics that are namespaced
    nicely."""

    def __init__(self, name):
        self.name_prefix = name

    def make_subspace(self, name):
        return Metrics("%s_%s" % (self.name_prefix, name))

    def register_collector(self, func):
        all_collectors.append(func)

    def _register(self, metric_class, name, *args, **kwargs):
        full_name = "%s_%s" % (self.name_prefix, name)

        metric = metric_class(full_name, *args, **kwargs)

        all_metrics.append(metric)
        return metric

    def register_counter(self, *args, **kwargs):
        return self._register(CounterMetric, *args, **kwargs)

    def register_callback(self, *args, **kwargs):
        return self._register(CallbackMetric, *args, **kwargs)

    def register_distribution(self, *args, **kwargs):
        return self._register(DistributionMetric, *args, **kwargs)

    def register_cache(self, *args, **kwargs):
        return self._register(CacheMetric, *args, **kwargs)

def register_memory_metrics(hs):
    try:
        import psutil
        process = psutil.Process()
        process.memory_info().rss
    except (ImportError, AttributeError):
        logger.warn(
            "psutil is not installed or incorrect version."
            " Disabling memory metrics."
        )
        return
    metric = MemoryUsageMetric(hs, psutil)
    all_metrics.append(metric)

def get_metrics_for(pkg_name):
    """ Returns a Metrics instance for conveniently creating metrics
    namespaced with the given name prefix. """

    # Convert a "package.name" to "package_name" because Prometheus doesn't
    # let us use . in metric names
    return Metrics(pkg_name.replace(".", "_"))

def render_all():
    strs = []

    for collector in all_collectors:
        collector()

    for metric in all_metrics:
        try:
            strs += metric.render()
        except Exception:
            strs += ["# FAILED to render"]
            logger.exception("Failed to render metric")

    strs.append("")  # to generate a final CRLF

    return "\n".join(strs)

reactor_metrics = get_metrics_for("reactor")
tick_time = reactor_metrics.register_distribution("tick_time")
pending_calls_metric = reactor_metrics.register_distribution("pending_calls")

gc_time = reactor_metrics.register_distribution("gc_time", labels=["gen"])
gc_unreachable = reactor_metrics.register_counter("gc_unreachable", labels=["gen"])

reactor_metrics.register_callback(
    "gc_counts", lambda: {(i,): v for i, v in enumerate(gc.get_count())}, labels=["gen"]
)

def runUntilCurrentTimer(func):

    @functools.wraps(func)
    def f(*args, **kwargs):
        now = reactor.seconds()
        num_pending = 0

        # _newTimedCalls is one long list of *all* pending calls. Below loop
        # is based off of impl of reactor.runUntilCurrent
        for delayed_call in reactor._newTimedCalls:
            if delayed_call.time > now:
                break

            if delayed_call.delayed_time > 0:
                continue

            num_pending += 1

        num_pending += len(reactor.threadCallQueue)

        start = time.time() * 1000
        ret = func(*args, **kwargs)
        end = time.time() * 1000
        tick_time.inc_by(end - start)
        pending_calls_metric.inc_by(num_pending)

        # Check if we need to do a manual GC (since its been disabled), and do
        # one if necessary.
        threshold = gc.get_threshold()
        counts = gc.get_count()
        for i in (2, 1, 0):
            if threshold[i] < counts[i]:
                logger.info("Collecting gc %d", i)

                start = time.time() * 1000
                unreachable = gc.collect(i)
                end = time.time() * 1000

                gc_time.inc_by(end - start, i)
                gc_unreachable.inc_by(unreachable, i)

        return ret

    return f

try:
    # Ensure the reactor has all the attributes we expect
    reactor.runUntilCurrent
    reactor._newTimedCalls
    reactor.threadCallQueue

    # runUntilCurrent is called when we have pending calls. It is called once
    # per iteratation after fd polling.
    reactor.runUntilCurrent = runUntilCurrentTimer(reactor.runUntilCurrent)

    # We manually run the GC each reactor tick so that we can get some metrics
    # about time spent doing GC,
    gc.disable()
except AttributeError:
    pass
<|endoftext|>


<commit_before>import discord
import asyncio
import re
import os
from discord.ext import commands
from .utils.dataIO import dataIO
from __main__ import send_cmd_help

class Dates:

    """Create dates for the chat"""

    def __init__(self, bot):
        self.bot = bot
        self.dates_path = "data/dates/dates.json"
        self.dates = dataIO.load_json(self.dates_path)

    @commands.group(name="termin", pass_context=True)
    async def date(self, ctx):
        """ Verwaltet Termine fürs Team"""
        if ctx.invoked_subcommand is None:
            await send_cmd_help(ctx)
    @date.command(name="aktivieren", pass_context=True)
    async def _activate(self, ctx):
        """Aktiviere die Terminverwaltung für diesen Server."""

        server = ctx.message.server
        if server.id not in self.dates:
            self.dates[server.id] = {}
            dataIO.save_json(self.dates_path, self.dates)

    @date.command(name="neu", pass_context=True)
    async def _new_date(self, ctx, date: str, time: str, note: str):
        """Erstelle einen neuen Termin."""

        newDate = (date, time, note)
        server = ctx.message.server
        datetime = date+time
        print(datetime)
        self.dates[server.id][datetime] = newDate
        dataIO.save_json(self.dates_path, self.dates)
        self.dates = dataIO.load_json(self.dates_path)

    @date.command(name="liste", pass_context=True)
    async def _list_dates(self, ctx):
        """Listet alle Termine auf."""

        server = ctx.message.server
        dates = self.dates[server.id]
        embed = discord.Embed(title="Termine")
        for date in dates:
            embed.add_field(text=self.dates[server.id][date])
            print(self.dates[server.id][date])
        await self.bot.say(embed=embed)

def check_folders():
    if not os.path.exists("data/dates"):
        print("Creating data/dates directory...")
        os.makedirs("data/dates")

def check_files():
    f = "data/dates/dates.json"
    if not dataIO.is_valid_json(f):
        print("Creating data/dates/dates.json...")
        dataIO.save_json(f, {})

def setup(bot):
    check_folders()
    check_files()
    n = Dates(bot)
    bot.add_cog(n)
<commit_msg>Fixed embed<commit_after>import discord
import asyncio
import re
import os
from discord.ext import commands
from .utils.dataIO import dataIO
from __main__ import send_cmd_help

class Dates:

    """Create dates for the chat"""

    def __init__(self, bot):
        self.bot = bot
        self.dates_path = "data/dates/dates.json"
        self.dates = dataIO.load_json(self.dates_path)

    @commands.group(name="termin", pass_context=True)
    async def date(self, ctx):
        """ Verwaltet Termine fürs Team"""
        if ctx.invoked_subcommand is None:
            await send_cmd_help(ctx)
    @date.command(name="aktivieren", pass_context=True)
    async def _activate(self, ctx):
        """Aktiviere die Terminverwaltung für diesen Server."""

        server = ctx.message.server
        if server.id not in self.dates:
            self.dates[server.id] = {}
            dataIO.save_json(self.dates_path, self.dates)

    @date.command(name="neu", pass_context=True)
    async def _new_date(self, ctx, date: str, time: str, note: str):
        """Erstelle einen neuen Termin."""

        newDate = (date, time, note)
        server = ctx.message.server
        datetime = date+time
        print(datetime)
        self.dates[server.id][datetime] = newDate
        dataIO.save_json(self.dates_path, self.dates)
        self.dates = dataIO.load_json(self.dates_path)

    @date.command(name="liste", pass_context=True)
    async def _list_dates(self, ctx):
        """Listet alle Termine auf."""

        server = ctx.message.server
        dates = self.dates[server.id]
        embed = discord.Embed(title="Termine")
        for date in dates:
            embed.add_field(name="Termin", value=self.dates[server.id][date])
            print(self.dates[server.id][date])
        await self.bot.say(embed=embed)

def check_folders():
    if not os.path.exists("data/dates"):
        print("Creating data/dates directory...")
        os.makedirs("data/dates")

def check_files():
    f = "data/dates/dates.json"
    if not dataIO.is_valid_json(f):
        print("Creating data/dates/dates.json...")
        dataIO.save_json(f, {})

def setup(bot):
    check_folders()
    check_files()
    n = Dates(bot)
    bot.add_cog(n)
<|endoftext|>