<commit_before>#! python
# -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function, division

import os
import sys
import time

from .parser import gsea_edb_parser, gsea_rank_metric, gsea_gmt_parser, gsea_cls_parser
from .algorithm import enrichment_score, gsea_compute, preprocess, ranking_metric
from .plot import gsea_plot, heatmap
from collections import OrderedDict

import pandas as pd

def replot(indir,outdir='gseapy_out', weight=1,figsize=[6.5,6], format='png',min_size=3, max_size=5000):
    """The main fuction to run inside python.

    :param indir: GSEA desktop results directory. In the sub folder, you must contain edb file foder.
    :param outdir: Output directory.
    :param weight: weighted score type. choose from {0,1,1.5,2}. Default: 1.
    :param figsize: matplotlib output figure figsize. Defult: [6.5,6].
    :param format: matplotlib output figure format. Default: 'png'.
    :param min_size: min size of input genes presented in Gene Sets. Default: 3.
    :param max_size: max size of input genes presented in Gene Sets. Default: 5000.
                     you will not encourage to use min_size, or max_size argment in :func:`replot` function.
                     Because gmt file has already been filter.

    :return: Generate new figures with seleted figure format. Default: 'png'.
    """
    argument = locals()
    import glob
    from bs4 import BeautifulSoup
    #parsing files.......
    results_path = glob.glob(indir+'*/edb/results.edb')[0]
    rank_path =  glob.glob(indir+'*/edb/*.rnk')[0]
    gene_set_path =  glob.glob(indir+'*/edb/gene_sets.gmt')[0]
    cls_path = glob.glob(indir+'*/edb/*.cls')[0]
    file_list = [results_path, rank_path, gene_set_path, cls_path]
    for file in file_list:
        if not os.path.isfile(file):
            print("Incorrect Input %s !" %file)
            sys.exit(1)
    #extract sample names from .cls file
    phenoPos, phenoNeg, classes = gsea_cls_parser(cls_path)
    #obtain gene sets
    gene_set_dict = gsea_gmt_parser(gene_set_path, min_size=min_size, max_size=max_size)
    #obtain rank_metrics
    rank_metric = gsea_rank_metric(rank_path)
    correl_vector =  rank_metric['rank'].values
    gene_list = rank_metric['gene_name']
    #extract each enriment term in the results.edb files and plot.
    database = BeautifulSoup(open(results_path),features='xml')
    length = len(database.findAll('DTG'))
    os.makedirs(outdir, exist_ok=True)

    for idx in range(length):
        #extract statistical resutls from results.edb file
        enrich_term, hit_ind, nes, pval, fdr= gsea_edb_parser(results_path, index=idx)
        gene_set = gene_set_dict.get(enrich_term)
        #calculate enrichment score
        RES = enrichment_score(gene_list=gene_list, gene_set=gene_set, weighted_score_type=weight,
                               correl_vector=correl_vector)[2]
        #plotting
        fig = gsea_plot(rank_metric, enrich_term,hit_ind, nes, pval,
                        fdr, RES, phenoPos, phenoNeg, figsize=figsize)
        fig.savefig('{a}/.gsea.replot.{b}.{c}'.format(a=outdir, b=enrich_term, c=format), dpi=300,)

    with open(outdir+"/command.txt",'wt') as f:
        argument = OrderedDict(sorted(argument.items(),key = lambda t:t[0]))
        for item in argument.items():
            f.write("%s = %s\n"%(item[0],item[1]))
    print("Congratulations! Your plots have been reproduced successfully!")

def call(data, gene_sets, cls, outdir='gseapy_out', min_size=15, max_size=1000, permutation_n=1000, weighted_score_type=1,
        permutation_type='gene_set', method='log2_ratio_of_classes', ascending=False, figsize=[6.5,6], format='png',
        graph_num=20, seed=None):
    """ Run Gene Set Enrichment Analysis.

    :param data: Gene expression data table.
    :param gene_sets: Gene sets file. e.g. gmt files. Same input with GSEA.
    :param permutation_n: Number of permutations for significance computation. Default: 1000.
    :param permutation_type: Permutation type, "phenotype" (default) for phenotypes, "gene_set" for genes.
    :param int min_size: Minimum allowed number of genes from gene set also the data set. Defaut: 15.
    :param int max_size: Maximum allowed number of genes from gene set also the data set. Defaults: 15 and 1000.
    :param weighted_score_type: Refer to :func:`algorithm.enrichment_socre`. Default:1.
    :param method: Ranking metric method, refer to :func:`algorithm.ranking_metric`.
    :param ascending: Sorting order of rankings. Default: False.
    :param outdir: Results output directory.
    :param figsize: Matplotlib figsize, accept a tuple or list, e.g. [width,height]. Default: [6.5,6].
    :param format: Matplotlib figure format. Default: 'png'.
    :param graph_num: Plot graphs for top sets of each phenotype
    :param seed: Random seed. expect an interger. Defalut:None.
    :return: Return a DataFrame when inside python console.
             Generate ``GSEA`` plots and store a dictionary into csv file,
             where dictionary key is a gene set and values are::

                 | {es: enrichment score,
                 |  nes: normalized enrichment score,
                 |  p: P-value,
                 |  fdr: FDR,
                 |  size: gene set size,
                 |  matched_size: genes matched to the data,
                 |  genes: gene names from the data set}

    """
    argument = locals()
    assert permutation_type in ["phenotype", "gene_set"]

    if isinstance(data, pd.DataFrame) :
        df = data.copy()
        argument['data'] = 'DataFrame'
    elif isinstance(data, str) :
        df = pd.read_table(data)
    else:
        raise Exception('Error parsing gene expression dataframe!')
        sys.exit(1)

    assert len(df) > 1

    #select correct expression genes and values.
    dat = preprocess(df)

    # phenotype labels parsing
    phenoPos, phenoNeg, classes = gsea_cls_parser(cls)

    #ranking metrics calculation.
    dat2 = ranking_metric(df=dat, method=method, phenoPos=phenoPos, phenoNeg=phenoNeg, classes=classes, ascending=ascending)

    #filtering out gene sets and build gene sets dictionary
    gmt = gsea_gmt_parser(gene_sets, min_size=min_size, max_size=max_size,gene_list=dat2['gene_name'].values)

    #compute ES, NES, pval, FDR, RES
    results,hit_ind,rank_ES, subsets = gsea_compute(data=dat, n=permutation_n,gmt=gmt, weighted_score_type=weighted_score_type,
                                                    permutation_type=permutation_type, method=method,
                                                    phenoPos=phenoPos, phenoNeg=phenoNeg, classes=classes, ascending=ascending,
                                                    seed=seed)

    os.makedirs(outdir, exist_ok=True)
    res = OrderedDict()
    for gs,gseale,ind,RES in zip(subsets, list(results), hit_ind, rank_ES):
        rdict = OrderedDict()
        rdict['es'] = gseale[0]
        rdict['nes'] = gseale[1]
        rdict['pval'] = gseale[2]
        rdict['fdr'] = gseale[3]
        rdict['gene_set_size'] = len(gmt[gs])
        rdict['matched_size'] = len(ind)
        rdict['rank_ES'] = RES
        rdict['genes'] = dat2.ix[ind,'gene_name'].tolist()
        rdict['hit_index'] = ind
        res[gs] = rdict

    res_df = pd.DataFrame.from_dict(res,orient='index')
    res_df.index.name = 'Term'
    res_df.sort_values(by='fdr', inplace=True)

    res_df.drop(['rank_ES','hit_index'], axis=1, inplace=True)
    res_df.to_csv('{a}/{b}.{c}.reports.csv'.format(a=outdir, b='gseapy', c=permutation_type), float_format ='%.7f')

    print("Start to generate gseapy reports, and produce figures...", time.ctime())

    #Plotting
    top_term = res_df.head(graph_num).index
    width = len(classes) if len(classes) >= 6 else  5
    for gs in top_term:
        hit = res.get(gs)['hit_index']
        gene_symbol = res.get(gs)['genes']
        fig = gsea_plot(rank_metric=dat2, enrich_term=gs, hit_ind=hit,
                        nes=res.get(gs)['nes'], pval=res.get(gs)['pval'], fdr=res.get(gs)['fdr'],
                        RES=res.get(gs)['rank_ES'], phenoPos=phenoPos, phenoNeg=phenoNeg, figsize=figsize)
        fig.savefig('{a}/{b}.gsea.{c}'.format(a=outdir, b=gs, c=format), dpi=300,)

        heatmap(df=dat.loc[gene_symbol], term=gs, outdir=outdir,
                figsize=(width, len(gene_symbol)/2), format=format)

    with open(outdir+"/command.txt",'wt') as f:
        argument = OrderedDict(sorted(argument.items(), key=lambda t:t[0]))
        for item in argument.items():
            f.write("%s = %s\n"%(item[0],item[1]))
    print("...Congratulations. GSEAPY run successfully!!!.............\n...The Job is done...........................Goodbye!")

    if isinstance(data, pd.DataFrame):
        return res_df

def prerank(rnk, gene_sets, outdir='gseapy_out', pheno_pos='Pos', pheno_neg='Neg',
            min_size=15, max_size=1000, permutation_n=1000, weighted_score_type=1,
            ascending=False, figsize=[6.5,6], format='png', graph_num=20, seed=None):
    """ Run Gene Set Enrichment Analysis with pre-ranked correlation defined by user.

    :param rnk: pre-ranked correlation table, Same input with ``GSEA`` .rnk file.
    :param gene_sets: Gene sets file. e.g. gmt files. Same input with GSEA.
    :param outdir: results output directory.
    :param permutation_n: Number of permutations for significance computation. Default: 1000.
    :param int min_size: Minimum allowed number of genes from gene set also the data set. Defaut: 15.
    :param int max_size: Maximum allowed number of genes from gene set also the data set. Defaults: 15 and 1000.
    :param weighted_score_type: Refer to :func:`algorithm.enrichment_socre`. Default:1.
    :param ascending: Sorting order of rankings. Default: False.
    :param figsize: Matplotlib figsize, accept a tuple or list, e.g. [width,height]. Default: [6.5,6].
    :param format: Matplotlib figure format. Default: 'png'.
    :param graph_num: Plot graphs for top sets of each phenotype
    :param seed: Random seed. expect an interger. Defalut:None.
    :return: Return a DataFrame when inside python console.
             Generate ``GSEA`` plots and store a dictionary into csv file,
             where dictionary key is a gene set and values are::

                 | {es: enrichment score,
                 |  nes: normalized enrichment score,
                 |  p: P-value,
                 |  fdr: FDR,
                 |  size: gene set size,
                 |  matched_size: genes matched to the data,
                 |  genes: gene names from the data set}

    """
    argument = locals()

    if isinstance(rnk, pd.DataFrame) :
        argument['rnk'] = 'DataFrame'

    dat2 = gsea_rank_metric(rnk)
    assert len(dat2) > 1
    #drop duplicates in ranking metrics.
    dat2.drop_duplicates(subset='gene_name',inplace=True, keep='first')
    assert len(dat2) > 1

    #filtering out gene sets and build gene sets dictionary
    gmt = gsea_gmt_parser(gene_sets, min_size=min_size, max_size=max_size, gene_list=dat2['gene_name'].values)

    #compute ES, NES, pval, FDR, RES
    results,hit_ind,rank_ES, subsets = gsea_compute(data=dat2, n=permutation_n, gmt=gmt, weighted_score_type=weighted_score_type,
                                                    permutation_type='gene_set', method=None, phenoPos=pheno_pos, phenoNeg=pheno_neg,
                                                    classes=None, ascending=ascending, seed=seed, prerank=True)

    print("Start to generate gseapy reports, and produce figures...", time.ctime())
    os.makedirs(outdir, exist_ok=True)
    res = OrderedDict()
    for gs,gseale,ind,RES in zip(subsets, list(results), hit_ind, rank_ES):
        rdict = OrderedDict()
        rdict['es'] = gseale[0]
        rdict['nes'] = gseale[1]
        rdict['pval'] = gseale[2]
        rdict['fdr'] = gseale[3]
        rdict['gene_set_size'] = len(gmt[gs])
        rdict['matched_size'] = len(ind)
        rdict['rank_ES'] = RES
        rdict['genes'] = dat2.ix[ind,'gene_name'].tolist()
        rdict['hit_index'] = ind
        res[gs] = rdict

    res_df = pd.DataFrame.from_dict(res, orient='index')
    res_df.index.name = 'Term'
    res_df.sort_values(by='fdr', inplace=True)

    res_df.drop(['rank_ES','hit_index'], axis=1, inplace=True)
    res_df.to_csv('{a}/{b}.prerank.reports.csv'.format(a=outdir, b='gseapy'), float_format ='%.7f')

    print("Start to generate gseapy reports, and produce figures...", time.ctime())

    #Plotting
    top_term = res_df.head(graph_num).index
    for gs in top_term:
        fig = gsea_plot(rank_metric=dat2, enrich_term=gs, hit_ind=res.get(gs)['hit_index'],
                        nes=res.get(gs)['nes'], pval=res.get(gs)['pval'], fdr=res.get(gs)['fdr'],
                        RES=res.get(gs)['rank_ES'], phenoPos=pheno_pos, phenoNeg=pheno_neg, figsize=figsize)
        fig.savefig('{a}/{b}.gsea.{c}'.format(a=outdir, b=gs, c=format), dpi=300,)

    with open(outdir+"/command.txt",'wt') as f:
        argument = OrderedDict(sorted(argument.items(),key = lambda t:t[0]))
        for item in argument.items():
            f.write("%s = %s\n"%(item[0], item[1]))
    print("Congratulations. GSEAPY run successfully................")
    print("The Job is done.................................Goodbye!", time.ctime())

    if isinstance(rnk, pd.DataFrame):
        return res_df
<commit_msg>max_size set to 500<commit_after>#! python
# -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function, division

import os
import sys
import time

from .parser import gsea_edb_parser, gsea_rank_metric, gsea_gmt_parser, gsea_cls_parser
from .algorithm import enrichment_score, gsea_compute, preprocess, ranking_metric
from .plot import gsea_plot, heatmap
from collections import OrderedDict

import pandas as pd

def replot(indir,outdir='gseapy_out', weight=1,figsize=[6.5,6], format='png',min_size=3, max_size=5000):
    """The main fuction to run inside python.

    :param indir: GSEA desktop results directory. In the sub folder, you must contain edb file foder.
    :param outdir: Output directory.
    :param weight: weighted score type. choose from {0,1,1.5,2}. Default: 1.
    :param figsize: matplotlib output figure figsize. Defult: [6.5,6].
    :param format: matplotlib output figure format. Default: 'png'.
    :param min_size: min size of input genes presented in Gene Sets. Default: 3.
    :param max_size: max size of input genes presented in Gene Sets. Default: 5000.
                     you will not encourage to use min_size, or max_size argment in :func:`replot` function.
                     Because gmt file has already been filter.

    :return: Generate new figures with seleted figure format. Default: 'png'.
    """
    argument = locals()
    import glob
    from bs4 import BeautifulSoup
    #parsing files.......
    results_path = glob.glob(indir+'*/edb/results.edb')[0]
    rank_path =  glob.glob(indir+'*/edb/*.rnk')[0]
    gene_set_path =  glob.glob(indir+'*/edb/gene_sets.gmt')[0]
    cls_path = glob.glob(indir+'*/edb/*.cls')[0]
    file_list = [results_path, rank_path, gene_set_path, cls_path]
    for file in file_list:
        if not os.path.isfile(file):
            print("Incorrect Input %s !" %file)
            sys.exit(1)
    #extract sample names from .cls file
    phenoPos, phenoNeg, classes = gsea_cls_parser(cls_path)
    #obtain gene sets
    gene_set_dict = gsea_gmt_parser(gene_set_path, min_size=min_size, max_size=max_size)
    #obtain rank_metrics
    rank_metric = gsea_rank_metric(rank_path)
    correl_vector =  rank_metric['rank'].values
    gene_list = rank_metric['gene_name']
    #extract each enriment term in the results.edb files and plot.
    database = BeautifulSoup(open(results_path),features='xml')
    length = len(database.findAll('DTG'))
    os.makedirs(outdir, exist_ok=True)

    for idx in range(length):
        #extract statistical resutls from results.edb file
        enrich_term, hit_ind, nes, pval, fdr= gsea_edb_parser(results_path, index=idx)
        gene_set = gene_set_dict.get(enrich_term)
        #calculate enrichment score
        RES = enrichment_score(gene_list=gene_list, gene_set=gene_set, weighted_score_type=weight,
                               correl_vector=correl_vector)[2]
        #plotting
        fig = gsea_plot(rank_metric, enrich_term,hit_ind, nes, pval,
                        fdr, RES, phenoPos, phenoNeg, figsize=figsize)
        fig.savefig('{a}/.gsea.replot.{b}.{c}'.format(a=outdir, b=enrich_term, c=format), dpi=300,)

    with open(outdir+"/command.txt",'wt') as f:
        argument = OrderedDict(sorted(argument.items(),key = lambda t:t[0]))
        for item in argument.items():
            f.write("%s = %s\n"%(item[0],item[1]))
    print("Congratulations! Your plots have been reproduced successfully!")

def call(data, gene_sets, cls, outdir='gseapy_out', min_size=15, max_size=500, permutation_n=1000, weighted_score_type=1,
        permutation_type='gene_set', method='log2_ratio_of_classes', ascending=False, figsize=[6.5,6], format='png',
        graph_num=20, seed=None):
    """ Run Gene Set Enrichment Analysis.

    :param data: Gene expression data table.
    :param gene_sets: Gene sets file. e.g. gmt files. Same input with GSEA.
    :param permutation_n: Number of permutations for significance computation. Default: 1000.
    :param permutation_type: Permutation type, "phenotype" (default) for phenotypes, "gene_set" for genes.
    :param int min_size: Minimum allowed number of genes from gene set also the data set. Defaut: 15.
    :param int max_size: Maximum allowed number of genes from gene set also the data set. Defaults: 500.
    :param weighted_score_type: Refer to :func:`algorithm.enrichment_socre`. Default:1.
    :param method: Ranking metric method, refer to :func:`algorithm.ranking_metric`.
    :param ascending: Sorting order of rankings. Default: False.
    :param outdir: Results output directory.
    :param figsize: Matplotlib figsize, accept a tuple or list, e.g. [width,height]. Default: [6.5,6].
    :param format: Matplotlib figure format. Default: 'png'.
    :param graph_num: Plot graphs for top sets of each phenotype
    :param seed: Random seed. expect an interger. Defalut:None.
    :return: Return a DataFrame when inside python console.
             Generate ``GSEA`` plots and store a dictionary into csv file,
             where dictionary key is a gene set and values are::

                 | {es: enrichment score,
                 |  nes: normalized enrichment score,
                 |  p: P-value,
                 |  fdr: FDR,
                 |  size: gene set size,
                 |  matched_size: genes matched to the data,
                 |  genes: gene names from the data set}

    """
    argument = locals()
    assert permutation_type in ["phenotype", "gene_set"]

    if isinstance(data, pd.DataFrame) :
        df = data.copy()
        argument['data'] = 'DataFrame'
    elif isinstance(data, str) :
        df = pd.read_table(data)
    else:
        raise Exception('Error parsing gene expression dataframe!')
        sys.exit(1)

    assert len(df) > 1

    #select correct expression genes and values.
    dat = preprocess(df)

    # phenotype labels parsing
    phenoPos, phenoNeg, classes = gsea_cls_parser(cls)

    #ranking metrics calculation.
    dat2 = ranking_metric(df=dat, method=method, phenoPos=phenoPos, phenoNeg=phenoNeg, classes=classes, ascending=ascending)

    #filtering out gene sets and build gene sets dictionary
    gmt = gsea_gmt_parser(gene_sets, min_size=min_size, max_size=max_size,gene_list=dat2['gene_name'].values)

    #compute ES, NES, pval, FDR, RES
    results,hit_ind,rank_ES, subsets = gsea_compute(data=dat, n=permutation_n,gmt=gmt, weighted_score_type=weighted_score_type,
                                                    permutation_type=permutation_type, method=method,
                                                    phenoPos=phenoPos, phenoNeg=phenoNeg, classes=classes, ascending=ascending,
                                                    seed=seed)

    os.makedirs(outdir, exist_ok=True)
    res = OrderedDict()
    for gs,gseale,ind,RES in zip(subsets, list(results), hit_ind, rank_ES):
        rdict = OrderedDict()
        rdict['es'] = gseale[0]
        rdict['nes'] = gseale[1]
        rdict['pval'] = gseale[2]
        rdict['fdr'] = gseale[3]
        rdict['gene_set_size'] = len(gmt[gs])
        rdict['matched_size'] = len(ind)
        rdict['rank_ES'] = RES
        rdict['genes'] = dat2.ix[ind,'gene_name'].tolist()
        rdict['hit_index'] = ind
        res[gs] = rdict

    res_df = pd.DataFrame.from_dict(res,orient='index')
    res_df.index.name = 'Term'
    res_df.sort_values(by='fdr', inplace=True)

    res_df.drop(['rank_ES','hit_index'], axis=1, inplace=True)
    res_df.to_csv('{a}/{b}.{c}.reports.csv'.format(a=outdir, b='gseapy', c=permutation_type), float_format ='%.7f')

    print("Start to generate gseapy reports, and produce figures...", time.ctime())

    #Plotting
    top_term = res_df.head(graph_num).index
    width = len(classes) if len(classes) >= 6 else  5
    for gs in top_term:
        hit = res.get(gs)['hit_index']
        gene_symbol = res.get(gs)['genes']
        fig = gsea_plot(rank_metric=dat2, enrich_term=gs, hit_ind=hit,
                        nes=res.get(gs)['nes'], pval=res.get(gs)['pval'], fdr=res.get(gs)['fdr'],
                        RES=res.get(gs)['rank_ES'], phenoPos=phenoPos, phenoNeg=phenoNeg, figsize=figsize)
        fig.savefig('{a}/{b}.gsea.{c}'.format(a=outdir, b=gs, c=format), dpi=300,)

        heatmap(df=dat.loc[gene_symbol], term=gs, outdir=outdir,
                figsize=(width, len(gene_symbol)/2), format=format)

    with open(outdir+"/command.txt",'wt') as f:
        argument = OrderedDict(sorted(argument.items(), key=lambda t:t[0]))
        for item in argument.items():
            f.write("%s = %s\n"%(item[0],item[1]))
    print("...Congratulations. GSEAPY run successfully!!!.............\n...The Job is done...........................Goodbye!")

    if isinstance(data, pd.DataFrame):
        return res_df

def prerank(rnk, gene_sets, outdir='gseapy_out', pheno_pos='Pos', pheno_neg='Neg',
            min_size=15, max_size=500, permutation_n=1000, weighted_score_type=1,
            ascending=False, figsize=[6.5,6], format='png', graph_num=20, seed=None):
    """ Run Gene Set Enrichment Analysis with pre-ranked correlation defined by user.

    :param rnk: pre-ranked correlation table, Same input with ``GSEA`` .rnk file.
    :param gene_sets: Gene sets file. e.g. gmt files. Same input with GSEA.
    :param outdir: results output directory.
    :param permutation_n: Number of permutations for significance computation. Default: 1000.
    :param int min_size: Minimum allowed number of genes from gene set also the data set. Defaut: 15.
    :param int max_size: Maximum allowed number of genes from gene set also the data set. Defaults: 500.
    :param weighted_score_type: Refer to :func:`algorithm.enrichment_socre`. Default:1.
    :param ascending: Sorting order of rankings. Default: False.
    :param figsize: Matplotlib figsize, accept a tuple or list, e.g. [width,height]. Default: [6.5,6].
    :param format: Matplotlib figure format. Default: 'png'.
    :param graph_num: Plot graphs for top sets of each phenotype
    :param seed: Random seed. expect an interger. Defalut:None.
    :return: Return a DataFrame when inside python console.
             Generate ``GSEA`` plots and store a dictionary into csv file,
             where dictionary key is a gene set and values are::

                 | {es: enrichment score,
                 |  nes: normalized enrichment score,
                 |  p: P-value,
                 |  fdr: FDR,
                 |  size: gene set size,
                 |  matched_size: genes matched to the data,
                 |  genes: gene names from the data set}

    """
    argument = locals()

    if isinstance(rnk, pd.DataFrame) :
        argument['rnk'] = 'DataFrame'

    dat2 = gsea_rank_metric(rnk)
    assert len(dat2) > 1
    #drop duplicates in ranking metrics.
    dat2.drop_duplicates(subset='gene_name',inplace=True, keep='first')
    assert len(dat2) > 1

    #filtering out gene sets and build gene sets dictionary
    gmt = gsea_gmt_parser(gene_sets, min_size=min_size, max_size=max_size, gene_list=dat2['gene_name'].values)

    #compute ES, NES, pval, FDR, RES
    results,hit_ind,rank_ES, subsets = gsea_compute(data=dat2, n=permutation_n, gmt=gmt, weighted_score_type=weighted_score_type,
                                                    permutation_type='gene_set', method=None, phenoPos=pheno_pos, phenoNeg=pheno_neg,
                                                    classes=None, ascending=ascending, seed=seed, prerank=True)

    print("Start to generate gseapy reports, and produce figures...", time.ctime())
    os.makedirs(outdir, exist_ok=True)
    res = OrderedDict()
    for gs,gseale,ind,RES in zip(subsets, list(results), hit_ind, rank_ES):
        rdict = OrderedDict()
        rdict['es'] = gseale[0]
        rdict['nes'] = gseale[1]
        rdict['pval'] = gseale[2]
        rdict['fdr'] = gseale[3]
        rdict['gene_set_size'] = len(gmt[gs])
        rdict['matched_size'] = len(ind)
        rdict['rank_ES'] = RES
        rdict['genes'] = dat2.ix[ind,'gene_name'].tolist()
        rdict['hit_index'] = ind
        res[gs] = rdict

    res_df = pd.DataFrame.from_dict(res, orient='index')
    res_df.index.name = 'Term'
    res_df.sort_values(by='fdr', inplace=True)

    res_df.drop(['rank_ES','hit_index'], axis=1, inplace=True)
    res_df.to_csv('{a}/{b}.prerank.reports.csv'.format(a=outdir, b='gseapy'), float_format ='%.7f')

    print("Start to generate gseapy reports, and produce figures...", time.ctime())

    #Plotting
    top_term = res_df.head(graph_num).index
    for gs in top_term:
        fig = gsea_plot(rank_metric=dat2, enrich_term=gs, hit_ind=res.get(gs)['hit_index'],
                        nes=res.get(gs)['nes'], pval=res.get(gs)['pval'], fdr=res.get(gs)['fdr'],
                        RES=res.get(gs)['rank_ES'], phenoPos=pheno_pos, phenoNeg=pheno_neg, figsize=figsize)
        fig.savefig('{a}/{b}.gsea.{c}'.format(a=outdir, b=gs, c=format), dpi=300,)

    with open(outdir+"/command.txt",'wt') as f:
        argument = OrderedDict(sorted(argument.items(),key = lambda t:t[0]))
        for item in argument.items():
            f.write("%s = %s\n"%(item[0], item[1]))
    print("Congratulations. GSEAPY run successfully................")
    print("The Job is done.................................Goodbye!", time.ctime())

    if isinstance(rnk, pd.DataFrame):
        return res_df
<|endoftext|>


<commit_before>import itertools
import collections
import typing as tp
import warnings

from ..recast_exceptions import rethrow_as, silence_excs
from ..decorators import for_argument

T, U = tp.TypeVar('T'), tp.TypeVar('U')
IteratorOrIterable = tp.Union[tp.Iterator[T], tp.Iterable[T]]

def length(iterator: IteratorOrIterable) -> int:
    """
    Return the length of an iterator, exhausting it by the way
    """
    i = 0
    for _ in iterator:
        i += 1
    return i

def walk(obj: T, child_getter: tp.Callable[[T], tp.Optional[tp.List[T]]] = list,
         deep_first: bool = True,
         leaves_only: bool = False) -> tp.Iterator[T]:
    """
    Return every node of a nested structure.

    :param obj: structure to traverse. This will not appear in generator
    :param child_getter: a callable to return a list of children of T.
        Should return an empty list or None of there are no more children.
    :param deep_first: if True, deep first will be returned, else it will be breadth first
    :param leaves_only: if True, only leaf nodes (having no children) will be returned
    """
    a = ConstruableIterator(child_getter(obj))
    for o in a:
        children = child_getter(o)
        if children is not None:
            try:
                child_len = len(children)
            except TypeError:
                child_len = 0

            if child_len:
                if deep_first:
                    a.add_many_immediate(children)
                else:
                    a.add_many(children)
                if leaves_only:
                    continue

        yield o

class ConstruableIterator:
    """
    An iterator that you can attach arbitrary things at the end and consume them during iteration.
    Eg:

    >>> a = ConstruableIterator([1, 2, 3])
    >>> for b in a:
    >>>   if b % 2 == 0:
    >>>       a.add(6)

    All arguments you provide to the constructor will be passed to underlying deque
    """
    __slots__ = ('entries', )

    def __init__(self, *args, **kwargs):
        self.entries = collections.deque(*args, **kwargs)       # type: tp.List[T]

    def add_immediate(self, t: T) -> None:
        """
        Schedule given value to be iterated over during the next __next__ call

        :param t: value to iterate over
        """
        self.entries.appendleft(t)

    def add_many_immediate(self, t: tp.Iterable[T]) -> None:
        """
        Schedule given values to be iterated over during the next __next__ call

        :param t: values to iterate over
        """
        for i, entry in enumerate(t):
            self.entries.insert(i, entry)

    def add(self, t: T) -> None:
        """
        Schedule given value to be iterated over after current items

        :param t: value to iterate over
        """
        self.entries.append(t)

    def add_many(self, t: tp.Iterable[T]) -> None:
        """
        Schedule given values to be iterated over after current items

        :param t: iterable of values
        """
        self.entries.extend(t)

    def __iter__(self) -> 'ConstruableIterator':
        return self

    @rethrow_as(IndexError, StopIteration)
    def __next__(self) -> T:
        return self.entries.popleft()

    def __length_hint__(self) -> int:
        return len(self.entries)

def unique(lst: IteratorOrIterable) -> tp.Iterator[T]:
    """
    Return each element from lst, but return every element only once.

    Take care for elements of T to be __eq__-able and hashable!

    This will keep internally a set of elements encountered, and skip them if same element
    appears twice

    :param lst: iterable to process
    :return: a generator yielding unique items from lst
    """
    already_seen = set()
    for elem in lst:
        if elem not in already_seen:
            already_seen.add(elem)
            yield elem

@for_argument(iter)
@silence_excs(StopIteration)
def even(sq: IteratorOrIterable) -> tp.Iterator[T]:
    """
    Return only elements with even indices in this iterable (first element will be returned,
    as indices are counted from 0)
    """
    while True:
        yield next(sq)
        next(sq)

@for_argument(iter)
@silence_excs(StopIteration)
def odd(sq: IteratorOrIterable) -> tp.Iterator[T]:
    """
    Return only elements with odd indices in this iterable.
    """
    while True:
        next(sq)
        yield next(sq)

def count(sq: IteratorOrIterable, start: tp.Optional[int] = None, step: int = 1,
          start_at: tp.Optional[int] = None) -> tp.Iterator[int]:
    """
    Return a sequence of integers, for each entry in the sequence with provided step.

    Essentially the same (if step were ignored) as:

    >>> (i for i, x in enumerate(sq, start=start_at))

    :param sq: sequence to enumerate
    :param start: alias for start_at. Prefer it in regards to start_at. Default is 0
    :param step: number to add to internal counter after each element
    :return: an iterator of subsequent integers
    """
    if start_at:
        start = start_at
    num = start or 0
    for _ in sq:
        yield num
        num += step

def iter_dict_of_list(dct: tp.Dict[T, tp.List[U]]) -> tp.Generator[tp.Tuple[T, U], None, None]:
    """
    Presents a simple way to iterate over a dictionary whose values are lists.

    This will return the dictionary key and each of the value contained in the list attached to
    the key.
    """
    for key, items in dct.items():
        for item in items:
            yield key, item

def infinite_counter(start_at: int = 0, step: int = 1) -> tp.Iterator[int]:
    """
    Infinite counter, starting at start_at

    :param start_at: value at which to start counting. It will be yielded as first
    :param step: step by which to progress the counter
    """
    warnings.warn('This is deprecated and will be removed in Satella 3.0, '
                  'use itertools.count() instead', DeprecationWarning)
    i = start_at
    while True:
        yield i
        i += step

def is_instance(classes: tp.Union[tp.Tuple[type, ...], type]) -> tp.Callable[[object], bool]:
    def inner(object_):
        return isinstance(object_, classes)

    inner.__doc__ = """Return a bool telling if object is of type %s""" % (repr(classes),)
    return inner

@for_argument(iter, iter)
@silence_excs(StopIteration)
def other_sequence_no_longer_than(base_sequence: IteratorOrIterable,
                                  other_sequence: IteratorOrIterable) -> tp.Iterator[T]:
    """
    Return every item in other_sequence, but limit it's p_len to that of base_sequence.

    If other_sequence is shorter than base_sequence, the shorter one will be returned.

    :param base_sequence: sequence whose p_len should be taken
    :param other_sequence: sequence to output values from
    """
    while True:
        next(base_sequence)
        yield next(other_sequence)

def shift(iterable_: IteratorOrIterable, shift_factor: int) -> tp.Iterator[T]:
    """
    Return this sequence, but shifted by factor elements, so that elements will appear
    sooner by factor.

    Eg:

    >>> assert list(shift([1,2, 3], 1)) == [2, 3, 1]

    However note that this will result in iterators which have negative shift to be readed entirely
    into memory (converted internally to lists). This can be avoided by passing in a Reversible
    iterable.

    :param iterable_: iterable to shift
    :param shift_factor: factor by which shift elements.
    :return: shifted sequence
    """

    if shift_factor >= 0:
        iterator = iter(iterable_)
        elements = []
        for i in range(shift_factor):
            elements.append(next(iterator))
        return itertools.chain(iterator, elements)
    else:
        if hasattr(iterable_, '__reversed__'):
            elements = take_n(reversed(iterable_), -shift_factor)
            elements = reversed(elements)
            return other_sequence_no_longer_than(iterable_, itertools.chain(elements, iterable_))
        else:
            iterator = list(iterable_)
            iterator = iterator[shift_factor:] + iterator[:shift_factor]  # shift's already negative
            return iterator

def zip_shifted(*args: tp.Union[IteratorOrIterable, tp.Tuple[IteratorOrIterable, int]]) -> \
        tp.Iterator[tp.Tuple[T, ...]]:
    """
    Construct an iterator, just like zip but first by cycling it's elements by it's shift factor.
    Elements will be shifted by a certain factor, this means that they will appear earlier.

    Example:

    >>> zip_shifted(([1, 2, 3, 4], 1), ([1, 2, 3, 4], 0)) == [(2, 1), (3, 2), (4, 3), (1, 4)]

    This will work on arbitrary iterators and iterables.

    Shift can be negative, in which case the last elements will appear sooner, eg.

    >>> zip_shifted(([1, 2, 3, 4], -1), ([1, 2, 3, 4], 0)) == [(4, 1), (1, 2), (2, 3), (3, 4)]

    Same memory considerations as :func:`shift` apply.

    The resulting iterator will be as long as the shortest sequence.
    :param args: a tuple with the iterator/iterable and amount of shift. If a non-tuple is given,
        it is assumed that the shift is zero.
    """
    warnings.warn('Use zip(shift(...)) instead!', DeprecationWarning)

    iterators = []  # type: tp.List[tp.Union[tp.Tuple[tp.Iterator[T], tp.List[T]], tp.Iterator[T]]
    for row in args:
        if not isinstance(row, tuple):
            iterators.append(row)
        else:
            iterable, shift_factor = row
            iterators.append(shift(iterable, shift_factor))
    return zip(*iterators)

@for_argument(iter)
@silence_excs(StopIteration)
def skip_first(iterator: IteratorOrIterable, n: int) -> tp.Iterator[T]:
    """
    Skip first n elements from given iterator.

    Returned iterator may be empty, if source iterator is shorter or equal to n.
    """
    warnings.warn('This is deprecated and will be removed in Satella 3.0. '
                  'Please use itertools.islice instead', DeprecationWarning)

    for i in range(n):
        next(iterator)
    yield from iterator

@for_argument(iter)
@silence_excs(StopIteration)
def stop_after(iterator: IteratorOrIterable, n: int) -> tp.Iterator[T]:
    """
    Stop this iterator after returning n elements, even if it's longer than that.

    The resulting iterator may be shorter than n, if the source element is so.

    :param iterator: iterator or iterable to examine
    :param n: elements to return
    """
    warnings.warn('This is deprecated and will be removed in Satella 3.0. '
                  'Please use itertools.islice instead', DeprecationWarning)

    for i in range(n):
        yield next(iterator)

@for_argument(iter)
def n_th(iterator: IteratorOrIterable, n: int = 0) -> T:
    """
    Obtain n-th element (counting from 0) of an iterable

    :param iterator: iterable to process
    :param n: element to return. Note that we're counting from 0
    :raises IndexError: iterable was too short
    """
    try:
        for i in range(n):
            next(iterator)
        return next(iterator)
    except (StopIteration, GeneratorExit):
        raise IndexError('Iterable was too short')

def map_list(fun: tp.Callable, iterable: IteratorOrIterable) -> tp.List:
    """
    A syntactic sugar for

    >>> list(map(fun, iterable))

    :param fun: function to apply
    :param iterable: iterable to iterate over
    """
    return list(map(fun, iterable))

def smart_enumerate(iterator: IteratorOrIterable, start: int = 0) -> tp.Iterator[tp.Tuple]:
    """
    An enumerate that talks pretty with lists of tuples. Consider

    >>> a = [(1, 2), (3, 4), (5, 6)]
    >>> for i, b in enumerate(a):
    >>>     c, d = b
    >>>     ...

    This function allows you just to write:
    >>> for i, c, d in enumerate(a):
    >>>     ...

    Note that elements in your iterable must be either a list of a tuple for that to work,
    or need to be able to be coerced to a tuple. Otherwise, TypeError will be thrown.

    :raise TypeError: could not coerce the elements in your iterable to a tuple
    """
    i = start
    for row in iterator:
        if isinstance(row, tuple):
            yield (i,) + row
        else:
            yield (i,) + tuple(row)
        i += 1

@for_argument(iter)
def take_n(iterator: IteratorOrIterable, n: int, skip: int = 0) -> tp.List[T]:
    """
    Take (first) n elements of an iterator, or the entire iterator, whichever comes first

    :param iterator: iterator to take from
    :param n: amount of elements to take
    :param skip: elements from the start to skip
    :return: list of p_len n (or shorter)
    """
    for i in range(skip):
        next(iterator)

    output = []
    for i in range(n):
        try:
            output.append(next(iterator))
        except StopIteration:
            return output
    return output
<commit_msg>2.11.1<commit_after>import itertools
import collections
import typing as tp
import warnings

from ..recast_exceptions import rethrow_as, silence_excs
from ..decorators import for_argument

T, U = tp.TypeVar('T'), tp.TypeVar('U')
IteratorOrIterable = tp.Union[tp.Iterator[T], tp.Iterable[T]]

def length(iterator: IteratorOrIterable) -> int:
    """
    Return the length of an iterator, exhausting it by the way
    """
    i = 0
    for _ in iterator:
        i += 1
    return i

def walk(obj: T, child_getter: tp.Callable[[T], tp.Optional[tp.List[T]]] = list,
         deep_first: bool = True,
         leaves_only: bool = False) -> tp.Iterator[T]:
    """
    Return every node of a nested structure.

    :param obj: structure to traverse. This will not appear in generator
    :param child_getter: a callable to return a list of children of T.
        Should return an empty list or None of there are no more children.
    :param deep_first: if True, deep first will be returned, else it will be breadth first
    :param leaves_only: if True, only leaf nodes (having no children) will be returned
    """
    a = ConstruableIterator(child_getter(obj))
    for o in a:
        children = child_getter(o)
        if children is not None:
            try:
                child_len = len(children)
            except TypeError:
                child_len = 0

            if child_len:
                if deep_first:
                    a.add_many_immediate(children)
                else:
                    a.add_many(children)
                if leaves_only:
                    continue

        yield o

class ConstruableIterator:
    """
    An iterator that you can attach arbitrary things at the end and consume them during iteration.
    Eg:

    >>> a = ConstruableIterator([1, 2, 3])
    >>> for b in a:
    >>>   if b % 2 == 0:
    >>>       a.add(6)

    All arguments you provide to the constructor will be passed to underlying deque
    """
    __slots__ = ('entries', )

    def __init__(self, *args, **kwargs):
        self.entries = collections.deque(*args, **kwargs)       # type: tp.List[T]

    def add_immediate(self, t: T) -> None:
        """
        Schedule given value to be iterated over during the next __next__ call

        :param t: value to iterate over
        """
        self.entries.appendleft(t)

    def add_many_immediate(self, t: tp.Iterable[T]) -> None:
        """
        Schedule given values to be iterated over during the next __next__ call

        :param t: values to iterate over
        """
        for i, entry in enumerate(t):
            self.entries.insert(i, entry)

    def add(self, t: T) -> None:
        """
        Schedule given value to be iterated over after current items

        :param t: value to iterate over
        """
        self.entries.append(t)

    def add_many(self, t: tp.Iterable[T]) -> None:
        """
        Schedule given values to be iterated over after current items

        :param t: iterable of values
        """
        self.entries.extend(t)

    def __iter__(self) -> 'ConstruableIterator':
        return self

    @rethrow_as(IndexError, StopIteration)
    def __next__(self) -> T:
        return self.entries.popleft()

    def __length_hint__(self) -> int:
        return len(self.entries)

def unique(lst: IteratorOrIterable) -> tp.Iterator[T]:
    """
    Return each element from lst, but return every element only once.

    Take care for elements of T to be __eq__-able and hashable!

    This will keep internally a set of elements encountered, and skip them if same element
    appears twice

    :param lst: iterable to process
    :return: a generator yielding unique items from lst
    """
    already_seen = set()
    for elem in lst:
        if elem not in already_seen:
            already_seen.add(elem)
            yield elem

@for_argument(iter)
def even(sq: IteratorOrIterable) -> tp.Iterator[T]:
    """
    Return only elements with even indices in this iterable (first element will be returned,
    as indices are counted from 0)
    """
    while True:
        try:
            yield next(sq)
            next(sq)
        except StopIteration:
            return

@for_argument(iter)
@silence_excs(StopIteration)
def odd(sq: IteratorOrIterable) -> tp.Iterator[T]:
    """
    Return only elements with odd indices in this iterable.
    """
    while True:
        try:
            next(sq)
            yield next(sq)
        except StopIteration:
            return

def count(sq: IteratorOrIterable, start: tp.Optional[int] = None, step: int = 1,
          start_at: tp.Optional[int] = None) -> tp.Iterator[int]:
    """
    Return a sequence of integers, for each entry in the sequence with provided step.

    Essentially the same (if step were ignored) as:

    >>> (i for i, x in enumerate(sq, start=start_at))

    :param sq: sequence to enumerate
    :param start: alias for start_at. Prefer it in regards to start_at. Default is 0
    :param step: number to add to internal counter after each element
    :return: an iterator of subsequent integers
    """
    if start_at:
        start = start_at
    num = start or 0
    for _ in sq:
        yield num
        num += step

def iter_dict_of_list(dct: tp.Dict[T, tp.List[U]]) -> tp.Generator[tp.Tuple[T, U], None, None]:
    """
    Presents a simple way to iterate over a dictionary whose values are lists.

    This will return the dictionary key and each of the value contained in the list attached to
    the key.
    """
    for key, items in dct.items():
        for item in items:
            yield key, item

def infinite_counter(start_at: int = 0, step: int = 1) -> tp.Iterator[int]:
    """
    Infinite counter, starting at start_at

    :param start_at: value at which to start counting. It will be yielded as first
    :param step: step by which to progress the counter
    """
    warnings.warn('This is deprecated and will be removed in Satella 3.0, '
                  'use itertools.count() instead', DeprecationWarning)
    i = start_at
    while True:
        yield i
        i += step

def is_instance(classes: tp.Union[tp.Tuple[type, ...], type]) -> tp.Callable[[object], bool]:
    def inner(object_):
        return isinstance(object_, classes)

    inner.__doc__ = """Return a bool telling if object is of type %s""" % (repr(classes),)
    return inner

@for_argument(iter, iter)
def other_sequence_no_longer_than(base_sequence: IteratorOrIterable,
                                  other_sequence: IteratorOrIterable) -> tp.Iterator[T]:
    """
    Return every item in other_sequence, but limit it's p_len to that of base_sequence.

    If other_sequence is shorter than base_sequence, the shorter one will be returned.

    :param base_sequence: sequence whose p_len should be taken
    :param other_sequence: sequence to output values from
    """
    while True:
        try:
            next(base_sequence)
            yield next(other_sequence)
        except StopIteration:
            return

def shift(iterable_: IteratorOrIterable, shift_factor: int) -> tp.Iterator[T]:
    """
    Return this sequence, but shifted by factor elements, so that elements will appear
    sooner by factor.

    Eg:

    >>> assert list(shift([1,2, 3], 1)) == [2, 3, 1]

    However note that this will result in iterators which have negative shift to be readed entirely
    into memory (converted internally to lists). This can be avoided by passing in a Reversible
    iterable.

    :param iterable_: iterable to shift
    :param shift_factor: factor by which shift elements.
    :return: shifted sequence
    """

    if shift_factor >= 0:
        iterator = iter(iterable_)
        elements = []
        for i in range(shift_factor):
            elements.append(next(iterator))
        return itertools.chain(iterator, elements)
    else:
        if hasattr(iterable_, '__reversed__'):
            elements = take_n(reversed(iterable_), -shift_factor)
            elements = reversed(elements)
            return other_sequence_no_longer_than(iterable_, itertools.chain(elements, iterable_))
        else:
            iterator = list(iterable_)
            iterator = iterator[shift_factor:] + iterator[:shift_factor]  # shift's already negative
            return iterator

@silence_excs(StopIteration)
def zip_shifted(*args: tp.Union[IteratorOrIterable, tp.Tuple[IteratorOrIterable, int]]) -> \
        tp.Iterator[tp.Tuple[T, ...]]:
    """
    Construct an iterator, just like zip but first by cycling it's elements by it's shift factor.
    Elements will be shifted by a certain factor, this means that they will appear earlier.

    Example:

    >>> zip_shifted(([1, 2, 3, 4], 1), ([1, 2, 3, 4], 0)) == [(2, 1), (3, 2), (4, 3), (1, 4)]

    This will work on arbitrary iterators and iterables.

    Shift can be negative, in which case the last elements will appear sooner, eg.

    >>> zip_shifted(([1, 2, 3, 4], -1), ([1, 2, 3, 4], 0)) == [(4, 1), (1, 2), (2, 3), (3, 4)]

    Same memory considerations as :func:`shift` apply.

    The resulting iterator will be as long as the shortest sequence.
    :param args: a tuple with the iterator/iterable and amount of shift. If a non-tuple is given,
        it is assumed that the shift is zero.
    """
    warnings.warn('This is deprecated and will be removed in Satella 3.0. '
                  'Use zip(shift(...)) instead!', DeprecationWarning)

    iterators = []  # type: tp.List[tp.Union[tp.Tuple[tp.Iterator[T], tp.List[T]], tp.Iterator[T]]
    for row in args:
        if not isinstance(row, tuple):
            iterators.append(row)
        else:
            iterable, shift_factor = row
            iterators.append(shift(iterable, shift_factor))
    return zip(*iterators)

@for_argument(iter)
@silence_excs(StopIteration)
def skip_first(iterator: IteratorOrIterable, n: int) -> tp.Iterator[T]:
    """
    Skip first n elements from given iterator.

    Returned iterator may be empty, if source iterator is shorter or equal to n.
    """
    warnings.warn('This is deprecated and will be removed in Satella 3.0. '
                  'Please use itertools.islice instead', DeprecationWarning)

    for i in range(n):
        next(iterator)
    yield from iterator

@for_argument(iter)
@silence_excs(StopIteration)
def stop_after(iterator: IteratorOrIterable, n: int) -> tp.Iterator[T]:
    """
    Stop this iterator after returning n elements, even if it's longer than that.

    The resulting iterator may be shorter than n, if the source element is so.

    :param iterator: iterator or iterable to examine
    :param n: elements to return
    """
    warnings.warn('This is deprecated and will be removed in Satella 3.0. '
                  'Please use itertools.islice instead', DeprecationWarning)

    for i in range(n):
        yield next(iterator)

@for_argument(iter)
def n_th(iterator: IteratorOrIterable, n: int = 0) -> T:
    """
    Obtain n-th element (counting from 0) of an iterable

    :param iterator: iterable to process
    :param n: element to return. Note that we're counting from 0
    :raises IndexError: iterable was too short
    """
    try:
        for i in range(n):
            next(iterator)
        return next(iterator)
    except (StopIteration, GeneratorExit):
        raise IndexError('Iterable was too short')

def map_list(fun: tp.Callable, iterable: IteratorOrIterable) -> tp.List:
    """
    A syntactic sugar for

    >>> list(map(fun, iterable))

    :param fun: function to apply
    :param iterable: iterable to iterate over
    """
    return list(map(fun, iterable))

def smart_enumerate(iterator: IteratorOrIterable, start: int = 0) -> tp.Iterator[tp.Tuple]:
    """
    An enumerate that talks pretty with lists of tuples. Consider

    >>> a = [(1, 2), (3, 4), (5, 6)]
    >>> for i, b in enumerate(a):
    >>>     c, d = b
    >>>     ...

    This function allows you just to write:
    >>> for i, c, d in enumerate(a):
    >>>     ...

    Note that elements in your iterable must be either a list of a tuple for that to work,
    or need to be able to be coerced to a tuple. Otherwise, TypeError will be thrown.

    :raise TypeError: could not coerce the elements in your iterable to a tuple
    """
    i = start
    for row in iterator:
        if isinstance(row, tuple):
            yield (i,) + row
        else:
            yield (i,) + tuple(row)
        i += 1

@for_argument(iter)
def take_n(iterator: IteratorOrIterable, n: int, skip: int = 0) -> tp.List[T]:
    """
    Take (first) n elements of an iterator, or the entire iterator, whichever comes first

    :param iterator: iterator to take from
    :param n: amount of elements to take
    :param skip: elements from the start to skip
    :return: list of p_len n (or shorter)
    """
    for i in range(skip):
        next(iterator)

    output = []
    for i in range(n):
        try:
            output.append(next(iterator))
        except StopIteration:
            return output
    return output
<|endoftext|>


<commit_before>#!/usr/bin/env python3

# gtk-example.py
# (c) Aleksander Alekseev 2016
# http://eax.me/

import signal
import os
import gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk
from gi.repository import Notify

APPID = "GTK Test"
CURRDIR = os.path.dirname(os.path.abspath(__file__))
# could be PNG or SVG as well
ICON = os.path.join(CURRDIR, 'python3.xpm')

# Cross-platform tray icon implementation
# See:
# * http://ubuntuforums.org/showthread.php?t=1923373#post11902222
# * https://github.com/syncthing/syncthing-gtk/blob/master/syncthing_gtk/statusicon.py

class TrayIcon:

    def __init__(self, appid, icon, menu):
        self.menu = menu

        APPIND_SUPPORT = 1
        try:
            from gi.repository import AppIndicator3
        except:
            APPIND_SUPPORT = 0

        if APPIND_SUPPORT == 1:
            self.ind = AppIndicator3.Indicator.new(
                appid, icon, AppIndicator3.IndicatorCategory.APPLICATION_STATUS)
            self.ind.set_status(AppIndicator3.IndicatorStatus.ACTIVE)
            self.ind.set_menu(self.menu)
        else:
            self.ind = Gtk.StatusIcon()
            self.ind.set_from_file(icon)
            self.ind.connect('popup-menu', self.right_click_event)

    def right_click_event(self, icon, button, time):
        def pos(menu, aicon):
            return (Gtk.StatusIcon.position_menu(menu, aicon))

        self.menu.popup(None, None, pos, icon, button, time)

class Handler:

    def __init__(self):
        self.window_is_hidden = False

    def onShowButtonClicked(self, button):
        msg = "Clicked: " + entry.get_text()
        dialog = Gtk.MessageDialog(window, 0, Gtk.MessageType.INFO,
                                   Gtk.ButtonsType.OK, msg)
        dialog.run()
        dialog.destroy()

    def onNotify(self, *args):
        Notify.Notification.new("Notification", "Hello!", ICON).show()

    def onShowOrHide(self, *args):
        if self.window_is_hidden:
            window.show()
        else:
            window.hide()

        self.window_is_hidden = not self.window_is_hidden

    def onQuit(self, *args):
        Notify.uninit()
        Gtk.main_quit()

# Handle pressing Ctr+C properly, ignored by default
signal.signal(signal.SIGINT, signal.SIG_DFL)

builder = Gtk.Builder()
builder.add_from_file('gtk-example.glade')
builder.connect_signals(Handler())

window = builder.get_object('window1')
window.set_icon_from_file(ICON)
window.show_all()

entry = builder.get_object('entry1')
menu = builder.get_object('menu1')
icon = TrayIcon(APPID, ICON, menu)
Notify.init(APPID)

Gtk.main()
<commit_msg>Minor changes for FreeBSD support<commit_after>#!/usr/bin/env python3

# gtk-example.py
# (c) Aleksander Alekseev 2016
# http://eax.me/

import signal
import os
import gi
gi.require_version('Gtk', '3.0')
gi.require_version('Notify', '0.7')
from gi.repository import Gtk
from gi.repository import Notify

APPID = "GTK Test"
CURRDIR = os.path.dirname(os.path.abspath(__file__))
# could be PNG or SVG as well
ICON = os.path.join(CURRDIR, 'python3.xpm')

# Cross-platform tray icon implementation
# See:
# * http://ubuntuforums.org/showthread.php?t=1923373#post11902222
# * https://github.com/syncthing/syncthing-gtk/blob/master/syncthing_gtk/statusicon.py
class TrayIcon:

    def __init__(self, appid, icon, menu):
        self.menu = menu

        APPIND_SUPPORT = 1
        try:
            from gi.repository import AppIndicator3
        except:
            APPIND_SUPPORT = 0

        if APPIND_SUPPORT == 1:
            self.ind = AppIndicator3.Indicator.new(
                appid, icon, AppIndicator3.IndicatorCategory.APPLICATION_STATUS)
            self.ind.set_status(AppIndicator3.IndicatorStatus.ACTIVE)
            self.ind.set_menu(self.menu)
        else:
            self.ind = Gtk.StatusIcon()
            self.ind.set_from_file(icon)
            self.ind.connect('popup-menu', self.right_click_event)

    def right_click_event(self, icon, button, time):
        def pos(menu, x, y, user_data):
            return (Gtk.StatusIcon.position_menu(menu, x, y, user_data))

        self.menu.popup(None, None, pos, icon, button, time)

class Handler:

    def __init__(self):
        self.window_is_hidden = False

    def onShowButtonClicked(self, button):
        msg = "Clicked: " + entry.get_text()
        dialog = Gtk.MessageDialog(window, 0, Gtk.MessageType.INFO,
                                   Gtk.ButtonsType.OK, msg)
        dialog.run()
        dialog.destroy()

    def onNotify(self, *args):
        Notify.Notification.new("Notification", "Hello!", ICON).show()

    def onShowOrHide(self, *args):
        if self.window_is_hidden:
            window.show()
        else:
            window.hide()

        self.window_is_hidden = not self.window_is_hidden

    def onQuit(self, *args):
        Notify.uninit()
        Gtk.main_quit()

# Handle pressing Ctr+C properly, ignored by default
signal.signal(signal.SIGINT, signal.SIG_DFL)

builder = Gtk.Builder()
builder.add_from_file('gtk-example.glade')
builder.connect_signals(Handler())

window = builder.get_object('window1')
window.set_icon_from_file(ICON)
window.show_all()

entry = builder.get_object('entry1')
menu = builder.get_object('menu1')
icon = TrayIcon(APPID, ICON, menu)
Notify.init(APPID)

Gtk.main()
<|endoftext|>


<commit_before>import six
from six import moves
from werkzeug import http

CRLF = b'\r\n'

def read_until(reader, delim):
    chunks = yield_until(reader, delim)
    return b''.join(chunks)

def read_headers(reader):
    headers = {}
    found = read_until_part(reader)
    for line in moves.filter(bool, found.split(CRLF)):
        name, value = line.decode('ascii').split(': ')
        name = name.lower()
        if name == 'content-disposition':
            for k, v in options(value):
                headers[k] = v
        else:
            headers[name] = value
    return headers

def options(value):
    _, options = http.parse_options_header(value)
    return six.iteritems(options)

def read_until_part(reader):
    delim = CRLF * 2
    data = read_until(reader, delim)
    return data

def parse_multipart(reader, boundary):
    read_until(reader, boundary)
    reader.read(len(boundary))

    while True:
        headers = read_headers(reader)
        reader.read(2 * len(CRLF))

        part = yield_until(reader, CRLF + boundary)
        yield headers, part

        reader.read(len(CRLF + boundary))
        end = reader.peek()
        if len(end) == 0 or end == b'--' + CRLF:
            break

def yield_until(reader, delim, size=65536):
    chunk = _peek(reader, 2 * size)
    while chunk:
        index = chunk.find(delim)
        if index == -1:
            end = min(len(chunk), size)
            yield chunk[:end]
            reader.read(end)
        else:
            yield chunk[:min(index, size)]
            if size < index:
                yield chunk[size:index]
            reader.read(index)
            break
        chunk = _peek(reader, 2 * size)

def _peek(reader, size):
    buf = reader.peek(size)
    return buf[:min(size, len(buf))]
<commit_msg>Use iter() for reading chunks<commit_after>import six
from six import moves
from werkzeug import http

CRLF = b'\r\n'

def read_until(reader, delim):
    chunks = yield_until(reader, delim)
    return b''.join(chunks)

def read_headers(reader):
    headers = {}
    found = read_until_part(reader)
    for line in moves.filter(bool, found.split(CRLF)):
        name, value = line.decode('ascii').split(': ')
        name = name.lower()
        if name == 'content-disposition':
            for k, v in options(value):
                headers[k] = v
        else:
            headers[name] = value
    return headers

def options(value):
    _, options = http.parse_options_header(value)
    return six.iteritems(options)

def read_until_part(reader):
    delim = CRLF * 2
    data = read_until(reader, delim)
    return data

def parse_multipart(reader, boundary):
    read_until(reader, boundary)
    reader.read(len(boundary))

    while True:
        headers = read_headers(reader)
        reader.read(2 * len(CRLF))

        part = yield_until(reader, CRLF + boundary)
        yield headers, part

        reader.read(len(CRLF + boundary))
        end = reader.peek()
        if len(end) == 0 or end == b'--' + CRLF:
            break

def yield_until(reader, delim, size=65536):
    for chunk in iter(lambda: _peek(reader, 2 * size), b''):
        index = chunk.find(delim)
        if index == -1:
            end = min(len(chunk), size)
            yield chunk[:end]
            reader.read(end)
        else:
            yield chunk[:min(index, size)]
            if size < index:
                yield chunk[size:index]
            reader.read(index)
            break

def _peek(reader, size):
    buf = reader.peek(size)
    return buf[:min(size, len(buf))]
<|endoftext|>


<commit_before>"""
Copyright (c) 2017 Fabian Affolter <fabian@affolter-engineering.ch>

Licensed under MIT. All rights reserved.
"""
import requests
import click

import pymystrom

URI = 'api/v1/device'
TIMEOUT = 5

@click.group()
@click.version_option()
def main():
    """Simple command-line tool to get and set the values of a myStrom devices.

    This tool can set the targets of a myStrom button for the different
    available actions single, double, long, and touch.
    """

@main.group('config')
def config():
    """Handler for all action around the configuration of a myStrom device."""

@config.command('read')
@click.option('--ip', prompt="IP address of the myStrom device",
              help="IP address of the myStrom device.")
@click.option('--mac', prompt="MAC address of the device",
              help="MAC address of the device.")
def read_config(ip, mac):
    """Read the current configuration of a myStrom device."""
    click.echo("Read configuration from %s" % ip)
    request = requests.get(
        'http://{}/{}/{}/'.format(ip, URI, mac), timeout=TIMEOUT)
    print(request.json())

@main.group('button')
def button():
    """Handler for all action around the configuration of a myStrom button."""

@button.command('generic')
@click.option('--ip', prompt="IP address of the myStrom button",
              help="IP address of the myStrom button.")
@click.option('--mac', prompt="MAC address of the button",
              help="MAC address of the myStrom button.")
@click.option('--single', prompt="URL for a single tap", default="",
              help="URL for a single tap.")
@click.option('--double', prompt="URL for a double tap", default="",
              help="URL for a double tap.")
@click.option('--long', prompt="URL for a long tab", default="",
              help="URL for a long tab.")
@click.option('--touch', prompt="URL for a touch", default="",
              help="URL for a touch.")
def write_config(ip, mac, single, double, long, touch):
    """Write the current configuration of a myStrom button."""
    click.echo("Write configuration to device %s" % ip)
    data =  {
        'single': single,
        'double': double,
        'long': long,
        'touch': touch,
    }
    request = requests.post(
        'http://{}/{}/{}/'.format(ip, URI, mac), data=data, timeout=TIMEOUT)

    if request.status_code == 200:
        click.echo("Configuration of %s set" % mac)

@button.command('home-assistant')
@click.option('--ip', prompt="IP address of the myStrom button",
              help="IP address of the myStrom button.")
@click.option('--mac', prompt="MAC address of the button",
              help="MAC address of the button.")
@click.option('--hass', prompt="IP address of the Home Assistant instance",
              help="IP address of Home Assistant instance to use.")
@click.option('--port', prompt="Port of Home Assistant instance",
              default="8123",
              help="Port where Home Assitant instance is listening.")
@click.option('--id', prompt="ID of the myStrom button", default="",
              help="ID of the myStrom button.")
def write_ha_config(ip, mac, hass, port, id):
    """Write the configuration for Home Assistant to a myStrom button."""
    click.echo("Write configuration for Home Assistant to device %s..." % ip)

    action = "get://{1}:{2}/api/mystrom?{0}={3}"
    data = {
        'single': action.format('single', hass, port, id),
        'double':  action.format('double', hass, port, id),
        'long':  action.format('long', hass, port, id),
        'touch':  action.format('touch', hass, port, id),
    }
    request = requests.post(
        'http://{}/{}/{}/'.format(ip, URI, mac), data=data, timeout=TIMEOUT)

    if request.status_code == 200:
        click.echo("Configuration for %s set" % ip)
        click.echo("After using the push pattern the first time then"
                   "the myStrom WiFi Button will show up as %s" % id)

@button.command('reset')
@click.option('--ip', prompt="IP address of the myStrom WiFi Button",
              help="P address of the myStrom WiFi Button.")
@click.option('--mac', prompt="MAC address of the myStrom button",
              help="MAC address of the myStrom Wifi Button.")
def reset_config(ip, mac):
    """Reset the current configuration of a myStrom WiFi Button."""
    click.echo("Reset configuration of button %s..." % ip)
    data = {
        'single': "",
        'double': "",
        'long': "",
        'touch': "",
    }
    request = requests.post(
        'http://{}/{}/{}/'.format(ip, URI, mac), data=data, timeout=TIMEOUT)

    if request.status_code == 200:
        click.echo("Reset configuration of %s" % mac)

@main.group('bulb')
def bulb():
    """Handler for all action around the configuration of a myStrom bulb."""

@bulb.command('on')
@click.option('--ip', prompt="IP address of the myStrom bulb",
              help="IP address of the myStrom bulb.")
@click.option('--mac', prompt="MAC address of the bulb",
              help="MAC address of the bulb.")
def on(ip, mac):
    """Switch the bulb on."""
    bulb = pymystrom.MyStromBulb(ip, mac)
    bulb.set_color_hex('000000FF')

@bulb.command('color')
@click.option('--ip', prompt="IP address of the myStrom bulb",
              help="IP address of the myStrom bulb.")
@click.option('--mac', prompt="MAC address of the bulb",
              help="MAC address of the bulb.")
@click.option('--hue', prompt="Set the hue of the bulb",
              help="Set the hue of the bulb.")
@click.option('--saturation', prompt="Set the saturation of the bulb",
              help="Set the saturation of the bulb.")
@click.option('--value', prompt="Set the value of the bulb",
              help="Set the value of the bulb.")
def color(ip, mac, hue, saturation, value):
    """Switch the bulb on with the given color."""
    bulb = pymystrom.MyStromBulb(ip, mac)
    bulb.set_color_hsv(hue, saturation, value)

@bulb.command('off')
@click.option('--ip', prompt="IP address of the myStrom bulb",
              help="IP address of the myStrom bulb.")
@click.option('--mac', prompt="MAC address of the bulb",
              help="MAC address of the bulb.")
def off(ip, mac):
    """Switch the bulb off."""
    bulb = pymystrom.MyStromBulb(ip, mac)
    bulb.set_off()

if __name__ == '__main__':
    main()
<commit_msg>Remove space<commit_after>"""
Copyright (c) 2017 Fabian Affolter <fabian@affolter-engineering.ch>

Licensed under MIT. All rights reserved.
"""
import requests
import click

import pymystrom

URI = 'api/v1/device'
TIMEOUT = 5

@click.group()
@click.version_option()
def main():
    """Simple command-line tool to get and set the values of a myStrom devices.

    This tool can set the targets of a myStrom button for the different
    available actions single, double, long, and touch.
    """

@main.group('config')
def config():
    """Handler for all action around the configuration of a myStrom device."""

@config.command('read')
@click.option('--ip', prompt="IP address of the myStrom device",
              help="IP address of the myStrom device.")
@click.option('--mac', prompt="MAC address of the device",
              help="MAC address of the device.")
def read_config(ip, mac):
    """Read the current configuration of a myStrom device."""
    click.echo("Read configuration from %s" % ip)
    request = requests.get(
        'http://{}/{}/{}/'.format(ip, URI, mac), timeout=TIMEOUT)
    print(request.json())

@main.group('button')
def button():
    """Handler for all action around the configuration of a myStrom button."""

@button.command('generic')
@click.option('--ip', prompt="IP address of the myStrom button",
              help="IP address of the myStrom button.")
@click.option('--mac', prompt="MAC address of the button",
              help="MAC address of the myStrom button.")
@click.option('--single', prompt="URL for a single tap", default="",
              help="URL for a single tap.")
@click.option('--double', prompt="URL for a double tap", default="",
              help="URL for a double tap.")
@click.option('--long', prompt="URL for a long tab", default="",
              help="URL for a long tab.")
@click.option('--touch', prompt="URL for a touch", default="",
              help="URL for a touch.")
def write_config(ip, mac, single, double, long, touch):
    """Write the current configuration of a myStrom button."""
    click.echo("Write configuration to device %s" % ip)
    data = {
        'single': single,
        'double': double,
        'long': long,
        'touch': touch,
    }
    request = requests.post(
        'http://{}/{}/{}/'.format(ip, URI, mac), data=data, timeout=TIMEOUT)

    if request.status_code == 200:
        click.echo("Configuration of %s set" % mac)

@button.command('home-assistant')
@click.option('--ip', prompt="IP address of the myStrom button",
              help="IP address of the myStrom button.")
@click.option('--mac', prompt="MAC address of the button",
              help="MAC address of the button.")
@click.option('--hass', prompt="IP address of the Home Assistant instance",
              help="IP address of Home Assistant instance to use.")
@click.option('--port', prompt="Port of Home Assistant instance",
              default="8123",
              help="Port where Home Assitant instance is listening.")
@click.option('--id', prompt="ID of the myStrom button", default="",
              help="ID of the myStrom button.")
def write_ha_config(ip, mac, hass, port, id):
    """Write the configuration for Home Assistant to a myStrom button."""
    click.echo("Write configuration for Home Assistant to device %s..." % ip)

    action = "get://{1}:{2}/api/mystrom?{0}={3}"
    data = {
        'single': action.format('single', hass, port, id),
        'double':  action.format('double', hass, port, id),
        'long':  action.format('long', hass, port, id),
        'touch':  action.format('touch', hass, port, id),
    }
    request = requests.post(
        'http://{}/{}/{}/'.format(ip, URI, mac), data=data, timeout=TIMEOUT)

    if request.status_code == 200:
        click.echo("Configuration for %s set" % ip)
        click.echo("After using the push pattern the first time then"
                   "the myStrom WiFi Button will show up as %s" % id)

@button.command('reset')
@click.option('--ip', prompt="IP address of the myStrom WiFi Button",
              help="P address of the myStrom WiFi Button.")
@click.option('--mac', prompt="MAC address of the myStrom button",
              help="MAC address of the myStrom Wifi Button.")
def reset_config(ip, mac):
    """Reset the current configuration of a myStrom WiFi Button."""
    click.echo("Reset configuration of button %s..." % ip)
    data = {
        'single': "",
        'double': "",
        'long': "",
        'touch': "",
    }
    request = requests.post(
        'http://{}/{}/{}/'.format(ip, URI, mac), data=data, timeout=TIMEOUT)

    if request.status_code == 200:
        click.echo("Reset configuration of %s" % mac)

@main.group('bulb')
def bulb():
    """Handler for all action around the configuration of a myStrom bulb."""

@bulb.command('on')
@click.option('--ip', prompt="IP address of the myStrom bulb",
              help="IP address of the myStrom bulb.")
@click.option('--mac', prompt="MAC address of the bulb",
              help="MAC address of the bulb.")
def on(ip, mac):
    """Switch the bulb on."""
    bulb = pymystrom.MyStromBulb(ip, mac)
    bulb.set_color_hex('000000FF')

@bulb.command('color')
@click.option('--ip', prompt="IP address of the myStrom bulb",
              help="IP address of the myStrom bulb.")
@click.option('--mac', prompt="MAC address of the bulb",
              help="MAC address of the bulb.")
@click.option('--hue', prompt="Set the hue of the bulb",
              help="Set the hue of the bulb.")
@click.option('--saturation', prompt="Set the saturation of the bulb",
              help="Set the saturation of the bulb.")
@click.option('--value', prompt="Set the value of the bulb",
              help="Set the value of the bulb.")
def color(ip, mac, hue, saturation, value):
    """Switch the bulb on with the given color."""
    bulb = pymystrom.MyStromBulb(ip, mac)
    bulb.set_color_hsv(hue, saturation, value)

@bulb.command('off')
@click.option('--ip', prompt="IP address of the myStrom bulb",
              help="IP address of the myStrom bulb.")
@click.option('--mac', prompt="MAC address of the bulb",
              help="MAC address of the bulb.")
def off(ip, mac):
    """Switch the bulb off."""
    bulb = pymystrom.MyStromBulb(ip, mac)
    bulb.set_off()

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