<commit_before># -*- coding: utf-8 -*-

import __main__ as main
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.transforms as transforms
from matplotlib.colors import Normalize
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure

class _MidpointNormalize(Normalize):
    def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):
        self.midpoint = midpoint
        Normalize.__init__(self, vmin, vmax, clip)

    def __call__(self, value, clip=None):
        # I'm ignoring masked values and all kinds of edge cases to make a
        # simple example...
        x, y = [self.vmin, self.midpoint, self.vmax], [0, 0.5, 1]
        return np.ma.masked_array(np.interp(value, x, y))

def z_score(data2d, axis=0):
    """Standarize the mean and variance of the data axis Parameters.

    :param data2d: DataFrame to normalize.
    :param axis: int, Which axis to normalize across. If 0, normalize across rows,
                  if 1, normalize across columns.

    :Returns: Normalized DataFrame. Noramlized data with a mean of 0 and variance of 1
              across the specified axis.

    """
    if axis == 1:
        z_scored = data2d
    else:
        z_scored = data2d.T

    z_scored = (z_scored - z_scored.mean()) / z_scored.std()

    if axis == 1:
        return z_scored
    else:
        return z_scored.T

def heatmap(df, term, outdir, axis=0, figsize=(5,5), format='png'):
    """Visualize the dataframe.

    :param df: DataFrame from expression table.
    :param term: gene set name.
    :param outdir: path to save heatmap.
    :param axis: z_score axis.
    :param figsize: heatmap figsize.
    :param format: Matplotlib supported figure formats.

    """
    df = z_score(df, axis=axis)

    # Get the positions and used label for the ticks
    nx, ny = df.T.shape
    xticks = np.arange(0, nx, 1) + .5
    yticks = np.arange(0, ny, 1) + .5

    fig = Figure(figsize=figsize)
    canvas = FigureCanvas(fig)
    ax = fig.add_subplot(111)
    vmin = np.percentile(df.min(), 2)
    vmax =  np.percentile(df.max(), 98)
    matrix = ax.pcolormesh(df.values, cmap=plt.cm.RdBu_r, vmin=vmin, vmax=vmax)
    ax.set_ylim([0,len(df)])
    ax.set(xticks=xticks, yticks=yticks)
    ax.set_xticklabels(df.columns.values, fontsize=18, rotation=90)
    ax.set_yticklabels(df.index.values,  fontsize=18)
    ax.set_title("%s\nHeatmap of the Analyzed GeneSet"%term, fontsize=24)
    ax.tick_params(axis='both', which='both', bottom='off', top='off',
                   right='off', left='off')

    #fig.colorbar(matrix, ax=ax)
    cax=fig.add_axes([0.93,0.25,0.05,0.20])
    cbar = fig.colorbar(matrix, cax=cax)
    cbar.ax.tick_params(axis='both', which='both', bottom='off', top='off',
                        right='off', left='off')
    for side in ["top", "right", "left", "bottom"]:
        ax.spines[side].set_visible(False)
        cbar.ax.spines[side].set_visible(False)
    #cbar.ax.set_title('',loc='left')

    canvas.print_figure("{a}/{b}.heatmap.{c}".format(a=outdir, b=term, c=format),
                        bbox_inches='tight')

def gsea_plot(rank_metric, enrich_term, hit_ind, nes, pval, fdr, RES,
              phenoPos=None, phenoNeg=None, figsize =(6.5,6), **kwarg):
    """This is the main function for reproducing the gsea plot.

    :param rank_metric: rankings, rank_metric['rank'].values.
    :param enrich_term: gene_set name
    :param hit_ind: hit indexs of rank_metric['gene_name'] presented in gene set S.
    :param nes: Normalized enrichment scores.
    :param pval: nominal p-value.
    :param fdr: false discoveray rate.
    :param RES: ranking enrichment scores of all genes in rank_metric['gene_name'].
    :param phenoPos: phenotype lable, positive correlated.
    :param phenoNeg: phenotype lable, negative correlated.
    :param figsize: matplotlib figsize.
    :return: fig object of gsea plot.
    """
    # plt.style.use('classic')
    # center color map at midpoint = 0
    norm = _MidpointNormalize(midpoint=0)

    #dataFrame of ranked matrix scores
    x = rank_metric.index.values
    #figsize = (6,6)
    phenoP_label = phenoPos + ' (Positively Correlated)'
    phenoN_label = phenoNeg + ' (Negatively Correlated)'
    zero_score_ind = np.abs(rank_metric['rank']).argmin()
    z_score_label = 'Zero score at ' + str(zero_score_ind)
    nes_label = 'NES: '+ "{:.3f}".format(float(nes))
    pval_label = 'Pval: '+ "{:.3f}".format(float(pval))
    fdr_label = 'FDR: '+ "{:.3f}".format(float(fdr))
    im_matrix = rank_metric.ix[:,1:].T

    #in most case, we will have mangy plots, so do not display plots
    #It's also convinient to run this script on command line.
    #plt.ioff()
    #GSEA Plots
    gs = plt.GridSpec(16,1)
    #fig = plt.figure(figsize=figsize)
    fig = Figure(figsize=figsize)
    canvas = FigureCanvas(fig)
    #Ranked Metric Scores Plot
    ax1 =  fig.add_subplot(gs[11:])
    ax1.fill_between(x, y1= rank_metric['rank'], y2=0, color='#C9D3DB')
    ax1.set_ylabel("Ranked list metric",fontsize=14)
    ax1.text(.05, .9, phenoP_label, color='red', horizontalalignment='left', verticalalignment='top',
         transform=ax1.transAxes)
    ax1.text(.95, .05, phenoN_label, color='Blue', horizontalalignment='right', verticalalignment='bottom',
         transform=ax1.transAxes)

    # the x coords of this transformation are data, and the y coord are axes
    trans1 = transforms.blended_transform_factory(ax1.transData, ax1.transAxes)
    ax1.vlines(zero_score_ind, 0, 1, linewidth=.5, transform=trans1, linestyles='--', color='grey')
    ax1.text(zero_score_ind, 0.5, z_score_label, horizontalalignment='center', verticalalignment='center',
             transform=trans1)
    ax1.set_xlabel("Rank in Ordered Dataset", fontsize=14)
    ax1.spines['top'].set_visible(False)
    ax1.tick_params(axis='both', which='both', top='off', right='off', left='off')
    ax1.locator_params(axis='y', nbins=5)
    ax1.yaxis.set_major_formatter(plt.FuncFormatter(lambda tick_loc,tick_num :  '{:.1f}'.format(tick_loc) ))

    # use round method to control float number
    #ax1.yaxis.set_major_formatter(plt.FuncFormatter(lambda tick_loc,tick_num :  round(tick_loc, 1) ))

    #gene hits
    ax2 = fig.add_subplot(gs[8:10], sharex=ax1)

    # the x coords of this transformation are data, and the y coord are axes
    trans2 = transforms.blended_transform_factory(ax2.transData, ax2.transAxes)
    ax2.vlines(hit_ind, 0, 1,linewidth=.5,transform=trans2)
    ax2.spines['bottom'].set_visible(False)
    ax2.tick_params(axis='both', which='both', bottom='off', top='off',
                    labelbottom='off', right='off', left='off',labelleft='off')
    #colormap
    ax3 =  fig.add_subplot(gs[10],sharex=ax1)
    ax3.imshow(im_matrix, aspect='auto', norm=norm, cmap=plt.cm.seismic, interpolation='none') # cm.coolwarm
    ax3.spines['bottom'].set_visible(False)
    ax3.tick_params(axis='both', which='both', bottom='off', top='off',
                    labelbottom='off', right='off', left='off',labelleft='off')

    # Enrichment score plot
    ax4 = fig.add_subplot(gs[:8],sharex=ax1)
    ax4.plot(x,RES,linewidth=4,color ='#88C544')
    ax4.text(.1, .1, fdr_label, transform=ax4.transAxes)
    ax4.text(.1, .2, pval_label, transform=ax4.transAxes)
    ax4.text(.1, .3, nes_label, transform=ax4.transAxes)

    # the y coords of this transformation are data, and the x coord are axes
    trans4 = transforms.blended_transform_factory(ax4.transAxes, ax4.transData)
    ax4.hlines(0, 0, 1, linewidth=.5, transform=trans4, color='grey')
    ax4.set_ylabel("Enrichment score (ES)", fontsize=14)
    ax4.set_xlim(min(x), max(x))
    ax4.tick_params(axis='both', which='both', bottom='off', top='off', labelbottom='off', right='off')
    ax4.locator_params(axis='y', nbins=5)
    # FuncFormatter need two argment, I don't know why. this lambda function used to format yaxis tick labels.
    ax4.yaxis.set_major_formatter(plt.FuncFormatter(lambda tick_loc,tick_num :  '{:.1f}'.format(tick_loc)) )

    #fig adjustment
    fig.suptitle(enrich_term, fontsize=16)
    fig.subplots_adjust(hspace=0)
    #fig.tight_layout()
    #plt.close(fig)

    return fig

def dotplot(df, cutoff=0.05, figsize=(3,6)):
    """Visualize enrichr or gsea results.

    :param df: GSEApy DataFrame results
    :return:  a dotplot for enrichr terms.
    """

    if 'fdr' in df.columns:
        #gsea results
        df.rename(columns={'pval':'Adjusted P-value',}, inplace=True)
        df['hits_ratio'] =  df['matched_size'] / df['gene_set_size']
    else:
        #enrichr results
        df['Count'] = df['Overlap'].str.split("/").str[0].astype(int)
        df['Background'] = df['Overlap'].str.split("/").str[1].astype(int)
        df['hits_ratio'] =  df['Count'] / df['Background']

    # pvalue cut off
    df = df[df['Adjusted P-value'] <= cutoff]

    if len(df) < 1:
        print("No enrich terms when cuttoff = %s"%cutoff )
        return None
    #sorting the dataframe for better visualization
    df = df.sort_values(by='Adjusted P-value', ascending=False)

    # x axis values
    padj = df['Adjusted P-value']
    x = - padj.apply(np.log10)
    # y axis index and values
    y=  [i for i in range(0,len(df))]
    labels = df.Term.values

    area = np.pi * (df['hits_ratio'] *50) **2

    #creat scatter plot
    if hasattr(main, '__file__'):
        fig = Figure(figsize=figsize)
        canvas = FigureCanvas(fig)
        ax = fig.add_subplot(111)
    else:
        fig, ax = plt.subplots(figsize=figsize)
    vmin = np.percentile(padj.min(), 2)
    vmax =  np.percentile(padj.max(), 98)
    sc = ax.scatter(x=x, y=y, s=area, edgecolors='face', c = padj,
                    cmap = plt.cm.RdBu,vmin=vmin, vmax=vmax)
    ax.set_xlabel("-log$_{10}$(Adjust P-value)")
    ax.yaxis.set_major_locator(plt.FixedLocator(y))
    ax.yaxis.set_major_formatter(plt.FixedFormatter(labels))
    ax.set_ylim([-1, len(df)])
    ax.grid()

    #colorbar
    cax=fig.add_axes([0.93,0.40,0.05,0.20])
    cbar = fig.colorbar(sc, cax=cax,)
    cbar.ax.tick_params(right='off')
    cbar.ax.set_title('Padj',loc='left')

    #scale of dots
    ax2 =fig.add_axes([0.93,0.55,0.05,0.12])

    #for terms less than 3
    if len(df) >= 3:
        x2 = [0]*3
        # find the index of the closest value to the median
        idx = [area.argmax(), np.abs(area - area.median()).argmin(), area.argmin()]
    else:
        x2 =  [0]*len(df)
        idx = df.index

    y2 = [i for i in range(0,len(x2))]
    ax2.scatter(x=x2, y=y2, s=area[idx], c='black', edgecolors='face')

    for i, index in enumerate(idx):
        ax2.text(x=0.8, y=y2[i], s=df['hits_ratio'][index].round(2),
                 verticalalignment='center', horizontalalignment='left')
    ax2.set_title("Gene\nRatio",loc='left')

    #turn off all spines and ticks
    ax2.axis('off')

    #plt.tight_layout()
    #canvas.print_figure('test', bbox_inches='tight')
    return fig

def adjust_spines(ax, spines):
    """function for removing spines and ticks.

    :param ax: axes object
    :param spines: a list of spines names to keep. e.g [left, right, top, bottom]
                    if spines = []. remove all spines and ticks.

    """
    for loc, spine in ax.spines.items():
        if loc in spines:
            spine.set_position(('outward', 10))  # outward by 10 points
            spine.set_smart_bounds(True)
        else:
            spine.set_color('none')  # don't draw spine

    # turn off ticks where there is no spine
    if 'left' in spines:
        ax.yaxis.set_ticks_position('left')
    else:
        # no yaxis ticks
        ax.yaxis.set_ticks([])

    if 'bottom' in spines:
        ax.xaxis.set_ticks_position('bottom')
    else:
        # no xaxis ticks
        ax.xaxis.set_ticks([])
<commit_msg>import order<commit_after># -*- coding: utf-8 -*-

import __main__ as main
import numpy as np

from matplotlib.colors import Normalize
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
import matplotlib.transforms as transforms

class _MidpointNormalize(Normalize):
    def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):
        self.midpoint = midpoint
        Normalize.__init__(self, vmin, vmax, clip)

    def __call__(self, value, clip=None):
        # I'm ignoring masked values and all kinds of edge cases to make a
        # simple example...
        x, y = [self.vmin, self.midpoint, self.vmax], [0, 0.5, 1]
        return np.ma.masked_array(np.interp(value, x, y))

def z_score(data2d, axis=0):
    """Standarize the mean and variance of the data axis Parameters.

    :param data2d: DataFrame to normalize.
    :param axis: int, Which axis to normalize across. If 0, normalize across rows,
                  if 1, normalize across columns.

    :Returns: Normalized DataFrame. Noramlized data with a mean of 0 and variance of 1
              across the specified axis.

    """
    if axis == 1:
        z_scored = data2d
    else:
        z_scored = data2d.T

    z_scored = (z_scored - z_scored.mean()) / z_scored.std()

    if axis == 1:
        return z_scored
    else:
        return z_scored.T

def heatmap(df, term, outdir, axis=0, figsize=(5,5), format='png'):
    """Visualize the dataframe.

    :param df: DataFrame from expression table.
    :param term: gene set name.
    :param outdir: path to save heatmap.
    :param axis: z_score axis.
    :param figsize: heatmap figsize.
    :param format: Matplotlib supported figure formats.

    """
    df = z_score(df, axis=axis)

    # Get the positions and used label for the ticks
    nx, ny = df.T.shape
    xticks = np.arange(0, nx, 1) + .5
    yticks = np.arange(0, ny, 1) + .5

    fig = Figure(figsize=figsize)
    canvas = FigureCanvas(fig)
    ax = fig.add_subplot(111)
    vmin = np.percentile(df.min(), 2)
    vmax =  np.percentile(df.max(), 98)
    matrix = ax.pcolormesh(df.values, cmap=plt.cm.RdBu_r, vmin=vmin, vmax=vmax)
    ax.set_ylim([0,len(df)])
    ax.set(xticks=xticks, yticks=yticks)
    ax.set_xticklabels(df.columns.values, fontsize=18, rotation=90)
    ax.set_yticklabels(df.index.values,  fontsize=18)
    ax.set_title("%s\nHeatmap of the Analyzed GeneSet"%term, fontsize=24)
    ax.tick_params(axis='both', which='both', bottom='off', top='off',
                   right='off', left='off')

    #fig.colorbar(matrix, ax=ax)
    cax=fig.add_axes([0.93,0.25,0.05,0.20])
    cbar = fig.colorbar(matrix, cax=cax)
    cbar.ax.tick_params(axis='both', which='both', bottom='off', top='off',
                        right='off', left='off')
    for side in ["top", "right", "left", "bottom"]:
        ax.spines[side].set_visible(False)
        cbar.ax.spines[side].set_visible(False)
    #cbar.ax.set_title('',loc='left')

    canvas.print_figure("{a}/{b}.heatmap.{c}".format(a=outdir, b=term, c=format),
                        bbox_inches='tight')

def gsea_plot(rank_metric, enrich_term, hit_ind, nes, pval, fdr, RES,
              phenoPos=None, phenoNeg=None, figsize =(6.5,6), **kwarg):
    """This is the main function for reproducing the gsea plot.

    :param rank_metric: rankings, rank_metric['rank'].values.
    :param enrich_term: gene_set name
    :param hit_ind: hit indexs of rank_metric['gene_name'] presented in gene set S.
    :param nes: Normalized enrichment scores.
    :param pval: nominal p-value.
    :param fdr: false discoveray rate.
    :param RES: ranking enrichment scores of all genes in rank_metric['gene_name'].
    :param phenoPos: phenotype lable, positive correlated.
    :param phenoNeg: phenotype lable, negative correlated.
    :param figsize: matplotlib figsize.
    :return: fig object of gsea plot.
    """
    # plt.style.use('classic')
    # center color map at midpoint = 0
    norm = _MidpointNormalize(midpoint=0)

    #dataFrame of ranked matrix scores
    x = rank_metric.index.values
    #figsize = (6,6)
    phenoP_label = phenoPos + ' (Positively Correlated)'
    phenoN_label = phenoNeg + ' (Negatively Correlated)'
    zero_score_ind = np.abs(rank_metric['rank']).argmin()
    z_score_label = 'Zero score at ' + str(zero_score_ind)
    nes_label = 'NES: '+ "{:.3f}".format(float(nes))
    pval_label = 'Pval: '+ "{:.3f}".format(float(pval))
    fdr_label = 'FDR: '+ "{:.3f}".format(float(fdr))
    im_matrix = rank_metric.ix[:,1:].T

    #in most case, we will have mangy plots, so do not display plots
    #It's also convinient to run this script on command line.
    #plt.ioff()
    #GSEA Plots
    gs = plt.GridSpec(16,1)
    #fig = plt.figure(figsize=figsize)
    fig = Figure(figsize=figsize)
    canvas = FigureCanvas(fig)
    #Ranked Metric Scores Plot
    ax1 =  fig.add_subplot(gs[11:])
    ax1.fill_between(x, y1= rank_metric['rank'], y2=0, color='#C9D3DB')
    ax1.set_ylabel("Ranked list metric",fontsize=14)
    ax1.text(.05, .9, phenoP_label, color='red', horizontalalignment='left', verticalalignment='top',
         transform=ax1.transAxes)
    ax1.text(.95, .05, phenoN_label, color='Blue', horizontalalignment='right', verticalalignment='bottom',
         transform=ax1.transAxes)

    # the x coords of this transformation are data, and the y coord are axes
    trans1 = transforms.blended_transform_factory(ax1.transData, ax1.transAxes)
    ax1.vlines(zero_score_ind, 0, 1, linewidth=.5, transform=trans1, linestyles='--', color='grey')
    ax1.text(zero_score_ind, 0.5, z_score_label, horizontalalignment='center', verticalalignment='center',
             transform=trans1)
    ax1.set_xlabel("Rank in Ordered Dataset", fontsize=14)
    ax1.spines['top'].set_visible(False)
    ax1.tick_params(axis='both', which='both', top='off', right='off', left='off')
    ax1.locator_params(axis='y', nbins=5)
    ax1.yaxis.set_major_formatter(plt.FuncFormatter(lambda tick_loc,tick_num :  '{:.1f}'.format(tick_loc) ))

    # use round method to control float number
    #ax1.yaxis.set_major_formatter(plt.FuncFormatter(lambda tick_loc,tick_num :  round(tick_loc, 1) ))

    #gene hits
    ax2 = fig.add_subplot(gs[8:10], sharex=ax1)

    # the x coords of this transformation are data, and the y coord are axes
    trans2 = transforms.blended_transform_factory(ax2.transData, ax2.transAxes)
    ax2.vlines(hit_ind, 0, 1,linewidth=.5,transform=trans2)
    ax2.spines['bottom'].set_visible(False)
    ax2.tick_params(axis='both', which='both', bottom='off', top='off',
                    labelbottom='off', right='off', left='off',labelleft='off')
    #colormap
    ax3 =  fig.add_subplot(gs[10],sharex=ax1)
    ax3.imshow(im_matrix, aspect='auto', norm=norm, cmap=plt.cm.seismic, interpolation='none') # cm.coolwarm
    ax3.spines['bottom'].set_visible(False)
    ax3.tick_params(axis='both', which='both', bottom='off', top='off',
                    labelbottom='off', right='off', left='off',labelleft='off')

    # Enrichment score plot
    ax4 = fig.add_subplot(gs[:8],sharex=ax1)
    ax4.plot(x,RES,linewidth=4,color ='#88C544')
    ax4.text(.1, .1, fdr_label, transform=ax4.transAxes)
    ax4.text(.1, .2, pval_label, transform=ax4.transAxes)
    ax4.text(.1, .3, nes_label, transform=ax4.transAxes)

    # the y coords of this transformation are data, and the x coord are axes
    trans4 = transforms.blended_transform_factory(ax4.transAxes, ax4.transData)
    ax4.hlines(0, 0, 1, linewidth=.5, transform=trans4, color='grey')
    ax4.set_ylabel("Enrichment score (ES)", fontsize=14)
    ax4.set_xlim(min(x), max(x))
    ax4.tick_params(axis='both', which='both', bottom='off', top='off', labelbottom='off', right='off')
    ax4.locator_params(axis='y', nbins=5)
    # FuncFormatter need two argment, I don't know why. this lambda function used to format yaxis tick labels.
    ax4.yaxis.set_major_formatter(plt.FuncFormatter(lambda tick_loc,tick_num :  '{:.1f}'.format(tick_loc)) )

    #fig adjustment
    fig.suptitle(enrich_term, fontsize=16)
    fig.subplots_adjust(hspace=0)
    #fig.tight_layout()
    #plt.close(fig)

    return fig

def dotplot(df, cutoff=0.05, figsize=(3,6)):
    """Visualize enrichr or gsea results.

    :param df: GSEApy DataFrame results
    :return:  a dotplot for enrichr terms.
    """

    if 'fdr' in df.columns:
        #gsea results
        df.rename(columns={'pval':'Adjusted P-value',}, inplace=True)
        df['hits_ratio'] =  df['matched_size'] / df['gene_set_size']
    else:
        #enrichr results
        df['Count'] = df['Overlap'].str.split("/").str[0].astype(int)
        df['Background'] = df['Overlap'].str.split("/").str[1].astype(int)
        df['hits_ratio'] =  df['Count'] / df['Background']

    # pvalue cut off
    df = df[df['Adjusted P-value'] <= cutoff]

    if len(df) < 1:
        print("No enrich terms when cuttoff = %s"%cutoff )
        return None
    #sorting the dataframe for better visualization
    df = df.sort_values(by='Adjusted P-value', ascending=False)

    # x axis values
    padj = df['Adjusted P-value']
    x = - padj.apply(np.log10)
    # y axis index and values
    y=  [i for i in range(0,len(df))]
    labels = df.Term.values

    area = np.pi * (df['hits_ratio'] *50) **2

    #creat scatter plot
    if hasattr(main, '__file__'):
        fig = Figure(figsize=figsize)
        canvas = FigureCanvas(fig)
        ax = fig.add_subplot(111)
    else:
        fig, ax = plt.subplots(figsize=figsize)
    vmin = np.percentile(padj.min(), 2)
    vmax =  np.percentile(padj.max(), 98)
    sc = ax.scatter(x=x, y=y, s=area, edgecolors='face', c = padj,
                    cmap = plt.cm.RdBu,vmin=vmin, vmax=vmax)
    ax.set_xlabel("-log$_{10}$(Adjust P-value)")
    ax.yaxis.set_major_locator(plt.FixedLocator(y))
    ax.yaxis.set_major_formatter(plt.FixedFormatter(labels))
    ax.set_ylim([-1, len(df)])
    ax.grid()

    #colorbar
    cax=fig.add_axes([0.93,0.40,0.05,0.20])
    cbar = fig.colorbar(sc, cax=cax,)
    cbar.ax.tick_params(right='off')
    cbar.ax.set_title('Padj',loc='left')

    #scale of dots
    ax2 =fig.add_axes([0.93,0.55,0.05,0.12])

    #for terms less than 3
    if len(df) >= 3:
        x2 = [0]*3
        # find the index of the closest value to the median
        idx = [area.argmax(), np.abs(area - area.median()).argmin(), area.argmin()]
    else:
        x2 =  [0]*len(df)
        idx = df.index

    y2 = [i for i in range(0,len(x2))]
    ax2.scatter(x=x2, y=y2, s=area[idx], c='black', edgecolors='face')

    for i, index in enumerate(idx):
        ax2.text(x=0.8, y=y2[i], s=df['hits_ratio'][index].round(2),
                 verticalalignment='center', horizontalalignment='left')
    ax2.set_title("Gene\nRatio",loc='left')

    #turn off all spines and ticks
    ax2.axis('off')

    #plt.tight_layout()
    #canvas.print_figure('test', bbox_inches='tight')
    return fig

def adjust_spines(ax, spines):
    """function for removing spines and ticks.

    :param ax: axes object
    :param spines: a list of spines names to keep. e.g [left, right, top, bottom]
                    if spines = []. remove all spines and ticks.

    """
    for loc, spine in ax.spines.items():
        if loc in spines:
            spine.set_position(('outward', 10))  # outward by 10 points
            spine.set_smart_bounds(True)
        else:
            spine.set_color('none')  # don't draw spine

    # turn off ticks where there is no spine
    if 'left' in spines:
        ax.yaxis.set_ticks_position('left')
    else:
        # no yaxis ticks
        ax.yaxis.set_ticks([])

    if 'bottom' in spines:
        ax.xaxis.set_ticks_position('bottom')
    else:
        # no xaxis ticks
        ax.xaxis.set_ticks([])
<|endoftext|>


<commit_before>import os

import pkg_resources

from ._version import __version__
# :obj:`str`: version

# API
from . import config
from . import core
from . import util
<commit_msg>removing unused import<commit_after>from ._version import __version__
# :obj:`str`: version

# API
from . import config
from . import core
from . import util
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
# =============================================================================
#      FileName: autoreload.py
#          Desc: get some referance from django
#        Author: ifkite
#         Email: holahello@163.com
#      HomePage: http://github.com/ifkite
#      python version: 2.7.10
#    LastChange: 2017-07-30 22:27:38

# =============================================================================
'''

import sys
import os
import signal
import threading
import time

def clean_filenames(filenames):
    """
    """
    filelist = []
    for filename in filenames:
        # if not filename:
        #    continue
        # https://stackoverflow.com/questions/8822335/what-do-the-python-file-extensions-pyc-pyd-pyo-stand-for
        if filename.endswith(".pyc") or filename.endswith(".pyo"):
            filename = filename[:-1]
        if os.path.exists(filename):
            filelist.append(filename)
    return filelist

def gen_filenames():
    modules = sys.modules.values()
    return clean_filenames([module.__file__ for module in modules if hasattr(module, '__file__')])

_mtimes = {}
FILE_CHANGED = 3
def is_file_changed():
    for filename in gen_filenames():
        stat = os.stat(filename)
        mtime = stat.st_mtime
        if filename not in _mtimes:
            _mtimes[filename] = mtime
            continue
        if mtime != _mtimes[filename]:
            _mtimes.clear()
            return True
    return False

def check_file_changed():
    while True:
        if is_file_changed():
            sys.exit(FILE_CHANGED)
            time.sleep(1)

def wait_child(pid):
    while True:
        wpid, sts = os.waitpid(pid, 0)
        if os.WIFSTOPPED(sts):
            continue
        # if receive keybord interuption or kill signal
        elif os.WIFSIGNALED(sts):
            return -os.WTERMSIG(sts)
        elif os.WIFEXITED(sts):
            return os.WEXITSTATUS(sts)
        else:
            raise "Not stopped, signaled or exited???"

def handle_child_exit(exit_code):
    if exit_code < 0:
        os.kill(os.getpid(), -exit_code)
    elif exit_code == FILE_CHANGED:
        pass
    else:
        sys.exit(exit_code)

def handle_parent_except(pid):
    os.kill(pid, signal.SIGKILL)
    sys.exit()

def restart_reloader():
    while True:
        args = [sys.executable] + sys.argv
        child_environ = os.environ.copy()
        pid = os.fork()

        # child process
        if not pid:
            child_environ["RUN_MAIN"] = "true"
            # may exit with FILE_CHANGED code
            # in fact, call itself
            os.execve(sys.executable, args, child_environ)

        # parent process
        else:
            try:
                exit_code = wait_child(pid)
                handle_child_exit(exit_code)
            # exceptions occured in parent process
            except KeyboardInterrupt:
                handle_parent_except(pid)

def autoreload(func, *args, **kwargs):
    # child process
    if os.environ.get("RUN_MAIN") == "true":
        thread = threading.Thread(target=func, args=args, kwargs=kwargs)
        thread.start()
        check_file_changed()

    # parent process
    else:
        restart_reloader()
<commit_msg>处理键盘中断<commit_after># -*- coding: utf-8 -*-
'''
# =============================================================================
#      FileName: autoreload.py
#          Desc: get some referance from django
#        Author: ifkite
#         Email: holahello@163.com
#      HomePage: http://github.com/ifkite
#      python version: 2.7.10
#    LastChange: 2017-07-30 22:27:38

# =============================================================================
'''

import sys
import os
import signal
import threading
import time

def clean_filenames(filenames):
    """
    """
    filelist = []
    for filename in filenames:
        # if not filename:
        #    continue
        # https://stackoverflow.com/questions/8822335/what-do-the-python-file-extensions-pyc-pyd-pyo-stand-for
        if filename.endswith(".pyc") or filename.endswith(".pyo"):
            filename = filename[:-1]
        if os.path.exists(filename):
            filelist.append(filename)
    return filelist

def gen_filenames():
    modules = sys.modules.values()
    return clean_filenames([module.__file__ for module in modules if hasattr(module, '__file__')])

_mtimes = {}
FILE_CHANGED = 3
def is_file_changed():
    for filename in gen_filenames():
        stat = os.stat(filename)
        mtime = stat.st_mtime
        if filename not in _mtimes:
            _mtimes[filename] = mtime
            continue
        if mtime != _mtimes[filename]:
            _mtimes.clear()
            return True
    return False

def check_file_changed():
    while True:
        if is_file_changed():
            sys.exit(FILE_CHANGED)
            time.sleep(1)

def wait_child(pid):
    while True:
        try:
            # wait for child process
            wpid, sts = os.waitpid(pid, 0)
        except KeyboardInterrupt:
            # handle exceptions when parent is waiting
            handle_parent_exit(pid)

        if os.WIFSTOPPED(sts):
            continue
        # if receive keybord interuption or kill signal
        elif os.WIFSIGNALED(sts):
            return -os.WTERMSIG(sts)
        elif os.WIFEXITED(sts):
            return os.WEXITSTATUS(sts)
        else:
            raise "Not stopped, signaled or exited???"

def handle_child_exit(exit_code):
    if exit_code < 0:
        os.kill(os.getpid(), -exit_code)
    elif exit_code == FILE_CHANGED:
        pass
    else:
        sys.exit(exit_code)

def handle_parent_exit(pid):
    os.kill(pid, signal.SIGKILL)
    sys.exit()

def restart_reloader():
    while True:
        args = [sys.executable] + sys.argv
        child_environ = os.environ.copy()
        pid = os.fork()

        # child process
        if not pid:
            child_environ["RUN_MAIN"] = "true"
            # may exit with FILE_CHANGED code
            # in fact, call itself
            os.execve(sys.executable, args, child_environ)

        # parent process
        else:
            exit_code = wait_child(pid)
            handle_child_exit(exit_code)

def autoreload(func, *args, **kwargs):
    # child process
    if os.environ.get("RUN_MAIN") == "true":
        thread = threading.Thread(target=func, args=args, kwargs=kwargs)
        thread.start()
        check_file_changed()

    # parent process
    else:
        restart_reloader()
<|endoftext|>


<commit_before>import ctypes
import pkg_resources
import os
import sys
from typing import List, Optional

from setuptools import setup

VERSION = '10.5.0'

# List of packages supported by this version of CuPy.
PACKAGES = [
    'cupy-cuda102',
    'cupy-cuda110',
    'cupy-cuda111',
    'cupy-cuda112',
    'cupy-cuda113',
    'cupy-cuda114',
    'cupy-cuda115',
    'cupy-cuda116',
    'cupy-rocm-4-0',
    'cupy-rocm-4-2',
    'cupy-rocm-4-3',
    'cupy-rocm-5-0',
]

# List of packages NOT supported by this version of CuPy.
PACKAGES_OUTDATED = [
    'cupy-cuda80',
    'cupy-cuda90',
    'cupy-cuda91',
    'cupy-cuda92',
    'cupy-cuda100',
    'cupy-cuda101',
]

# List of sdist packages.
PACKAGES_SDIST = [
    'cupy',
]

class AutoDetectionFailed(Exception):
    pass

def _log(msg: str) -> None:
    sys.stdout.write(f'[cupy-wheel] {msg}\n')
    sys.stdout.flush()

def _get_version_from_library(
        libnames: List[str],
        funcname: str,
        nvrtc: bool = False,
) -> Optional[int]:
    """Returns the library version from list of candidate libraries."""

    for libname in libnames:
        try:
            _log(f'Looking for library: {libname}')
            runtime_so = ctypes.CDLL(libname)
            break
        except Exception as e:
            _log(f'Failed to open {libname}: {e}')
    else:
        _log('No more candidate library to find')
        return None

    func = getattr(runtime_so, funcname, None)
    if func is None:
        raise AutoDetectionFailed(
            f'{libname}: {func} could not be found')
    func.restype = ctypes.c_int

    if nvrtc:
        # nvrtcVersion
        func.argtypes = [
            ctypes.POINTER(ctypes.c_int),
            ctypes.POINTER(ctypes.c_int),
        ]
        major = ctypes.c_int()
        minor = ctypes.c_int()
        retval = func(major, minor)
        version = major.value * 1000 + minor.value * 10
    else:
        # cudaRuntimeGetVersion
        func.argtypes = [
            ctypes.POINTER(ctypes.c_int),
        ]
        version_ref = ctypes.c_int()
        retval = func(version_ref)
        version = version_ref.value

    if retval != 0:  # NVRTC_SUCCESS or cudaSuccess
        raise AutoDetectionFailed(
            f'{libname}: {func} returned error: {retval}')
    _log(f'Detected version: {version}')
    return version

def _setup_win32_dll_directory() -> None:
    if not hasattr(os, 'add_dll_directory'):
        # Python 3.7 or earlier.
        return
    cuda_path = os.environ.get('CUDA_PATH', None)
    if cuda_path is None:
        _log('CUDA_PATH is not set.'
             'cupy-wheel may not be able to discover NVRTC to probe version')
        return
    os.add_dll_directory(os.path.join(cuda_path, 'bin'))  # type: ignore[attr-defined] # NOQA

def _get_cuda_version() -> Optional[int]:
    """Returns the detected CUDA version or None."""

    if sys.platform == 'linux':
        libnames = [
            'libnvrtc.so.11.2',
            'libnvrtc.so.11.1',
            'libnvrtc.so.11.0',
            'libnvrtc.so.10.2',
        ]
    elif sys.platform == 'win32':
        libnames = [
            'nvrtc64_112_0.dll',
            'nvrtc64_111_0.dll',
            'nvrtc64_110_0.dll',
            'nvrtc64_102_0.dll',
        ]
        _setup_win32_dll_directory()
    else:
        _log(f'CUDA detection unsupported on platform: {sys.platform}')
        return None
    _log(f'Trying to detect CUDA version from libraries: {libnames}')
    version = _get_version_from_library(libnames, 'nvrtcVersion', True)
    return version

def _get_rocm_version() -> Optional[int]:
    """Returns the detected ROCm version or None."""
    if sys.platform == 'linux':
        libnames = ['libamdhip64.so']
    else:
        _log(f'ROCm detection unsupported on platform: {sys.platform}')
        return None
    version = _get_version_from_library(libnames, 'hipRuntimeGetVersion')
    return version

def _find_installed_packages() -> List[str]:
    """Returns the list of CuPy packages installed in the environment."""

    found = []
    for pkg in (PACKAGES + PACKAGES_OUTDATED + PACKAGES_SDIST):
        try:
            pkg_resources.get_distribution(pkg)
            found.append(pkg)
        except pkg_resources.DistributionNotFound:
            pass
    return found

def _cuda_version_to_package(ver: int) -> str:
    if ver < 10020:
        raise AutoDetectionFailed(
            f'Your CUDA version ({ver}) is too old.')
    elif ver < 11000:
        # CUDA 10.2
        suffix = '102'
    elif ver < 11010:
        # CUDA 11.0
        suffix = '110'
    elif ver < 11020:
        # CUDA 11.1
        suffix = '111'
    elif ver < 11030:
        # CUDA 11.2
        suffix = '112'
    elif ver < 11040:
        # CUDA 11.3
        suffix = '113'
    elif ver < 11050:
        # CUDA 11.4
        suffix = '114'
    elif ver < 11060:
        # CUDA 11.5
        suffix = '115'
    elif ver < 11070:
        # CUDA 11.6
        suffix = '116'
    else:
        raise AutoDetectionFailed(
            f'Your CUDA version ({ver}) is too new.')
    return f'cupy-cuda{suffix}'

def _rocm_version_to_package(ver: int) -> str:
    """
    ROCm 4.0.x = 3212
    ROCm 4.1.x = 3241
    ROCm 4.2.0 = 3275
    ROCm 4.3.0 = 40321300
    ROCm 4.3.1 = 40321331
    ROCm 4.5.0 = 40421401
    ROCm 4.5.1 = 40421432
    ROCm 5.0.0 = 50013601
    ROCm 5.1.0 = 50120531
    """
    if ver == 3212:
        # ROCm 4.0
        suffix = '4-0'
    elif ver == 3275:
        # ROCm 4.2
        suffix = '4-2'
    elif 4_03_00000 <= ver < 4_04_00000:
        # ROCm 4.3
        suffix = '4-3'
    elif 5_00_00000 <= ver < 5_01_00000:
        # ROCm 5.0
        suffix = '5-0'
    else:
        raise AutoDetectionFailed(
            f'Your ROCm version ({ver}) is unsupported.')
    return f'cupy-rocm-{suffix}'

def infer_best_package() -> str:
    """Returns the appropriate CuPy wheel package name for the environment."""

    # Find the existing CuPy wheel installation.
    installed = _find_installed_packages()
    if 1 < len(installed):
        raise AutoDetectionFailed(
            'You have multiple CuPy packages installed: \n'
            f'  {installed}\n'
            'Please uninstall all of them first, then try reinstalling.')

    elif 1 == len(installed):
        if installed[0] in PACKAGES_SDIST:
            raise AutoDetectionFailed(
                'You already have CuPy installed via source'
                ' (pip install cupy).')
        if installed[0] in PACKAGES_OUTDATED:
            raise AutoDetectionFailed(
                f'You have CuPy package "{installed[0]}" installed, but the'
                f' package is not available for version {VERSION}.')
        return installed[0]

    # Try CUDA.
    version = _get_cuda_version()
    if version is not None:
        return _cuda_version_to_package(version)

    # Try ROCm.
    version = _get_rocm_version()
    if version is not None:
        return _rocm_version_to_package(version)

    raise AutoDetectionFailed(
        'Unable to detect NVIDIA CUDA or AMD ROCm installation.')

#
# Entrypoint
#

def main() -> None:
    if os.environ.get('CUPY_UNIVERSAL_PKG_BUILD', None) is None:
        package = infer_best_package()
        requires = f'{package}=={VERSION}'
        _log(f'Installing package: {requires}')
        install_requires = [requires]
    else:
        _log('Building cupy-wheel package for release.')
        install_requires = []

    setup(
        name='cupy-wheel',
        version=VERSION,
        install_requires=install_requires,
    )

if __name__ == '__main__':
    main()
<commit_msg>cupy-wheel 10.6.0<commit_after>import ctypes
import pkg_resources
import os
import sys
from typing import List, Optional

from setuptools import setup

VERSION = '10.6.0'

# List of packages supported by this version of CuPy.
PACKAGES = [
    'cupy-cuda102',
    'cupy-cuda110',
    'cupy-cuda111',
    'cupy-cuda112',
    'cupy-cuda113',
    'cupy-cuda114',
    'cupy-cuda115',
    'cupy-cuda116',
    'cupy-cuda117',
    'cupy-rocm-4-0',
    'cupy-rocm-4-2',
    'cupy-rocm-4-3',
    'cupy-rocm-5-0',
]

# List of packages NOT supported by this version of CuPy.
PACKAGES_OUTDATED = [
    'cupy-cuda80',
    'cupy-cuda90',
    'cupy-cuda91',
    'cupy-cuda92',
    'cupy-cuda100',
    'cupy-cuda101',
]

# List of sdist packages.
PACKAGES_SDIST = [
    'cupy',
]

class AutoDetectionFailed(Exception):
    pass

def _log(msg: str) -> None:
    sys.stdout.write(f'[cupy-wheel] {msg}\n')
    sys.stdout.flush()

def _get_version_from_library(
        libnames: List[str],
        funcname: str,
        nvrtc: bool = False,
) -> Optional[int]:
    """Returns the library version from list of candidate libraries."""

    for libname in libnames:
        try:
            _log(f'Looking for library: {libname}')
            runtime_so = ctypes.CDLL(libname)
            break
        except Exception as e:
            _log(f'Failed to open {libname}: {e}')
    else:
        _log('No more candidate library to find')
        return None

    func = getattr(runtime_so, funcname, None)
    if func is None:
        raise AutoDetectionFailed(
            f'{libname}: {func} could not be found')
    func.restype = ctypes.c_int

    if nvrtc:
        # nvrtcVersion
        func.argtypes = [
            ctypes.POINTER(ctypes.c_int),
            ctypes.POINTER(ctypes.c_int),
        ]
        major = ctypes.c_int()
        minor = ctypes.c_int()
        retval = func(major, minor)
        version = major.value * 1000 + minor.value * 10
    else:
        # cudaRuntimeGetVersion
        func.argtypes = [
            ctypes.POINTER(ctypes.c_int),
        ]
        version_ref = ctypes.c_int()
        retval = func(version_ref)
        version = version_ref.value

    if retval != 0:  # NVRTC_SUCCESS or cudaSuccess
        raise AutoDetectionFailed(
            f'{libname}: {func} returned error: {retval}')
    _log(f'Detected version: {version}')
    return version

def _setup_win32_dll_directory() -> None:
    if not hasattr(os, 'add_dll_directory'):
        # Python 3.7 or earlier.
        return
    cuda_path = os.environ.get('CUDA_PATH', None)
    if cuda_path is None:
        _log('CUDA_PATH is not set.'
             'cupy-wheel may not be able to discover NVRTC to probe version')
        return
    os.add_dll_directory(os.path.join(cuda_path, 'bin'))  # type: ignore[attr-defined] # NOQA

def _get_cuda_version() -> Optional[int]:
    """Returns the detected CUDA version or None."""

    if sys.platform == 'linux':
        libnames = [
            'libnvrtc.so.11.2',
            'libnvrtc.so.11.1',
            'libnvrtc.so.11.0',
            'libnvrtc.so.10.2',
        ]
    elif sys.platform == 'win32':
        libnames = [
            'nvrtc64_112_0.dll',
            'nvrtc64_111_0.dll',
            'nvrtc64_110_0.dll',
            'nvrtc64_102_0.dll',
        ]
        _setup_win32_dll_directory()
    else:
        _log(f'CUDA detection unsupported on platform: {sys.platform}')
        return None
    _log(f'Trying to detect CUDA version from libraries: {libnames}')
    version = _get_version_from_library(libnames, 'nvrtcVersion', True)
    return version

def _get_rocm_version() -> Optional[int]:
    """Returns the detected ROCm version or None."""
    if sys.platform == 'linux':
        libnames = ['libamdhip64.so']
    else:
        _log(f'ROCm detection unsupported on platform: {sys.platform}')
        return None
    version = _get_version_from_library(libnames, 'hipRuntimeGetVersion')
    return version

def _find_installed_packages() -> List[str]:
    """Returns the list of CuPy packages installed in the environment."""

    found = []
    for pkg in (PACKAGES + PACKAGES_OUTDATED + PACKAGES_SDIST):
        try:
            pkg_resources.get_distribution(pkg)
            found.append(pkg)
        except pkg_resources.DistributionNotFound:
            pass
    return found

def _cuda_version_to_package(ver: int) -> str:
    if ver < 10020:
        raise AutoDetectionFailed(
            f'Your CUDA version ({ver}) is too old.')
    elif ver < 11000:
        # CUDA 10.2
        suffix = '102'
    elif ver < 11010:
        # CUDA 11.0
        suffix = '110'
    elif ver < 11020:
        # CUDA 11.1
        suffix = '111'
    elif ver < 11030:
        # CUDA 11.2
        suffix = '112'
    elif ver < 11040:
        # CUDA 11.3
        suffix = '113'
    elif ver < 11050:
        # CUDA 11.4
        suffix = '114'
    elif ver < 11060:
        # CUDA 11.5
        suffix = '115'
    elif ver < 11070:
        # CUDA 11.6
        suffix = '116'
    elif ver < 11080:
        # CUDA 11.7
        suffix = '117'
    else:
        raise AutoDetectionFailed(
            f'Your CUDA version ({ver}) is too new.')
    return f'cupy-cuda{suffix}'

def _rocm_version_to_package(ver: int) -> str:
    """
    ROCm 4.0.x = 3212
    ROCm 4.1.x = 3241
    ROCm 4.2.0 = 3275
    ROCm 4.3.0 = 40321300
    ROCm 4.3.1 = 40321331
    ROCm 4.5.0 = 40421401
    ROCm 4.5.1 = 40421432
    ROCm 5.0.0 = 50013601
    ROCm 5.1.0 = 50120531
    """
    if ver == 3212:
        # ROCm 4.0
        suffix = '4-0'
    elif ver == 3275:
        # ROCm 4.2
        suffix = '4-2'
    elif 4_03_00000 <= ver < 4_04_00000:
        # ROCm 4.3
        suffix = '4-3'
    elif 5_00_00000 <= ver < 5_01_00000:
        # ROCm 5.0
        suffix = '5-0'
    else:
        raise AutoDetectionFailed(
            f'Your ROCm version ({ver}) is unsupported.')
    return f'cupy-rocm-{suffix}'

def infer_best_package() -> str:
    """Returns the appropriate CuPy wheel package name for the environment."""

    # Find the existing CuPy wheel installation.
    installed = _find_installed_packages()
    if 1 < len(installed):
        raise AutoDetectionFailed(
            'You have multiple CuPy packages installed: \n'
            f'  {installed}\n'
            'Please uninstall all of them first, then try reinstalling.')

    elif 1 == len(installed):
        if installed[0] in PACKAGES_SDIST:
            raise AutoDetectionFailed(
                'You already have CuPy installed via source'
                ' (pip install cupy).')
        if installed[0] in PACKAGES_OUTDATED:
            raise AutoDetectionFailed(
                f'You have CuPy package "{installed[0]}" installed, but the'
                f' package is not available for version {VERSION}.')
        return installed[0]

    # Try CUDA.
    version = _get_cuda_version()
    if version is not None:
        return _cuda_version_to_package(version)

    # Try ROCm.
    version = _get_rocm_version()
    if version is not None:
        return _rocm_version_to_package(version)

    raise AutoDetectionFailed(
        'Unable to detect NVIDIA CUDA or AMD ROCm installation.')

#
# Entrypoint
#

def main() -> None:
    if os.environ.get('CUPY_UNIVERSAL_PKG_BUILD', None) is None:
        package = infer_best_package()
        requires = f'{package}=={VERSION}'
        _log(f'Installing package: {requires}')
        install_requires = [requires]
    else:
        _log('Building cupy-wheel package for release.')
        install_requires = []

    setup(
        name='cupy-wheel',
        version=VERSION,
        install_requires=install_requires,
    )

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


<commit_before>#!/usr/bin/python

import common
#-----
import re
import os
import string
import sys
import time
import traceback
import ConfigParser
import StringIO

import inspect

#### Class Configuration ############################
class Configuration:

  def __init__ (self, fileName):
    self.cp       = None
    self.inifile = fileName

    #-- check for duplicate sections in config file --
    self.check_for_duplicate_sections()

    try:
      #-- check version of parser to use ---
      if sys.version_info[0] >= 2 and sys.version_info[1] >= 3:
        self.cp = ConfigParser.SafeConfigParser()
      else:
        self.cp = ConfigParser.ConfigParser()
      #-- read the ini file ---
      real_fp = open(self.inifile, 'r')
      string_fp = StringIO.StringIO(real_fp.read())
      self.cp.readfp(string_fp,self.inifile)
    except Exception, e:
      common.logerr("%s" % e)

    #-- additional check for syntax errors --
    self.syntax_check()

  #----------------
  def syntax_check(self):
    """ Checks for some syntax errors in ini config file. """
    for section in self.sections():
      for option in self.options(section):
        value = self.option_value(section,option)
        if "\n" in value:
          line = string.split(value,"\n")
          common.logerr("Section [%s]: this line starts with whitespace ( %s)\n       Please remove the leading space or comment (;) the line." % (section,line[1]))

  #----------------
  def __str__ (self):
    result = []
    result.append('<Configuration from %s>' % self.inifile)
    for section in self.sections():
      result.append('[%s]' % section)
      for option in self.options(section):
        value = self.cp.get(section, option)
        result.append('    %-25s %s' % (option, value))
    return '\n'.join(result)

  #----------------
  def check_for_duplicate_sections(self):
    """ Check for duplicate sections in a config file ignoring commented ones.
        In addition, checks to verify there is no whitespace preceding or
        appending the section name in brackets as the ini parser does not
        validate for this.
    """
    if (self.inifile == "") or (self.inifile is None):
      common.logerr("System error: config file name is empty")
    try:
      fp = open(self.inifile, 'r')
    except:
      common.logerr("Problem reading ini file: %s" % sys.exc_info()[1])
    sections = {}    # used to keep track of sections
    duplicates = []  # used to identify duplicates
    for line in fp.readlines():
      newline = line.strip()
      if len(newline) == 0:
        continue
      if newline[0] != "[":
        continue
      #--- see if it is a section ---
      match = re.search('\[(.*)\]', newline)
      if match:
        section = match.group(1).lower().strip()
        if section in sections:
          duplicates.append(section)
          continue
        sections[section] = True
    if (len(duplicates) <> 0 ):
      common.logerr("Duplicate sections in %s - %s" % (self.inifile,duplicates))

  #----------------
  def validate_section(self,section,valid_option_list):
    if not self.has_section(section):
      common.logerr("Section (%s) does not exist in ini file (%s)" % (section,self.inifile))
    errors = []
    for option in valid_option_list:
      if self.has_option(section,option):
        continue
      errors.append(option)
    if len(errors) > 0:
      common.logerr("These options are not defined in the %s section of the ini file: %s" % (section,errors))

  #----------------
  def section_options(self):
    result = []
    for section in self.sections():
      result.append('[%s]' % section)
      for option in self.options(section):
        result.append('    %-25s' % (option))
    return '\n'.join(result)

  #----------------
  def filename(self):
    return self.inifile

  #----------------
  def sections(self):
    sections = list(self.cp.sections())
    sections.sort()
    return sections

  #----------------
  def options(self,section):
    options = self.cp.options(section)
    options.sort()
    return options

  #----------------
  def option_value(self,section,option):
    return self.cp.get(section,option)

  #----------------
  def has_section(self, section):
    return self.cp.has_section(section)
  #----------------
  def has_option(self, section,option):
    return self.cp.has_option(section,option)
  #----------------
  def delete_section(self, section):
    self.cp.remove_section(section)
    return

#### Exceptions #####################################
class ConfigurationError(Exception):
  pass
class UsageError(Exception):
  pass

#####################################################
#---------------------
#---- Functions ------
#---------------------
def compare_ini_files (file_1, file_2, no_local_settings):
  try:
    print "Comparing ini files: %s / %s" % (file_1,file_2)
    ini_1 = Configuration(file_1)
    ini_2 = Configuration(file_2)
    rtn = 0
    #--- remove the Local Settings section conditionally ---
    if ( no_local_settings ):
      ini_1.delete_section("Local Settings")
      ini_2.delete_section("Local Settings")
    print "... Checking section information:"
    if ( ini_1.sections() == ini_2.sections() ):
      print "... sections are identical"
    else:
      print "... WARNING: section information differs"
      compare_sections(ini_1,ini_2)
      compare_sections(ini_2,ini_1)
      rtn = 1
    print
    print "... Checking section/object information:"
    if ( ini_1.section_options() ==  ini_2.section_options() ):
      print "... all section/objects are identical"
    else:
      print "... WARNING: section/object information differs"
      compare_options(ini_1,ini_2)
      compare_options(ini_2,ini_1)
      rtn = 1
  except:
    raise
  return rtn

#--------------------------------
def compare_sections(ini1,ini2):
  print """
... Sections     in %s
       NOT FOUND in %s""" % (ini1.filename(),ini2.filename())
  for section in ini1.sections():
    if ( ini2.has_section(section) ):
      continue
    else:
      print "    %s" % (section)

#--------------------------------
def compare_options(ini1,ini2):
  print """
... Section/objects in %s
          NOT FOUND in %s""" % (ini1.filename(),ini2.filename())
  for section in ini1.sections():
    for option in ini1.options(section):
      ## if (section == "SE CHANGEME"):
      ##   if (option == "enable"):
      ##     print section,option
      if ( ini2.has_option(section,option) == False):
        print "    %-20s/%s" % (section,option)

#--------------------------------
def usage(pgm):
  print
  print "Usage: " + pgm + " --compare file1 file2"
  print "       " + pgm + " --show-options file"
  print "       " + pgm + " --show-values file"
  print "       " + pgm + " --validate file"
  print "       " + pgm + " --help | -h "
  print """
   compare .......... Shows the differences between the
                      section/objects (not values) of the 2 ini files
                        returns 0 if identical
                        returns 1 if any differences
   show-options ...... Shows the section/objects for the ini file
   show-values ....... Shows the section/objects/values for the ini file
   validate .......... Verifies the ini file has no syntax errors

   Full path to the files must be specified unless this is executed
   in the directory in which they reside.

"""

#----------------------------------
def run_unit_tests(pgm):
  try:
    dir="./testdata/"
    tests = {
"no arguments"        : [1,pgm],
"not enough arguments": [1,pgm,"--validate"],
"not enough arguments": [1,pgm,"--compare"],
"invalid argument"    : [1,pgm,"--bad-arg",dir+"non-existent-file"],

##  validate ###
"validate: good ini"  : [0,pgm,"--validate",dir+"config-good.ini"],
"validate: bad ini"   : [1,pgm,"--validate",dir+"config-bad.ini"],
"validate: no ini"    : [1,pgm,"--validate",dir+"non-existent-file"],
"duplicates"          : [1,pgm,"--validate",dir+"config-w-dup-sections.ini"],

##  compare ###
"compare: no difference" : [0,pgm,"--compare",dir+"config-good.ini",dir+"config-good.ini"],
"compare: differences"   : [1,pgm,"--compare",dir+"config-good.ini",dir+"config-good-different.ini"],
"compare: no ini"        : [1,pgm,"--validate",dir+"non-existent-file"],
"duplicates"             : [1,pgm,"--compare",dir+"config-good.ini config-w-dup-sections.ini"],

##  show-options ###
"show-options"          : [0,pgm,"--show-options",dir+"config-good.ini"],
"show-options: bad ini" : [1,pgm,"--show-options",dir+"config-bad.ini"],
"show-options: no ini"  : [1,pgm,"--show-options",dir+"non-existent-file"],
"duplicates"            : [1,pgm,"--show-options",dir+"config-w-dup-sections.ini"],

##  show-values ###
"show-values"           : [0,pgm,"--show-values",dir+"config-good.ini"],
"show-values: bad ini"  : [1,pgm,"--show-values",dir+"config-bad.ini"],
"show-values: no ini"   : [1,pgm,"--show-values",dir+"non-existent-file"],
"duplicates"            : [1,pgm,"--show-values",dir+"config-w-dup-sections.ini"],
    }
    #---- run tests -----
    n=0
    for test in tests.keys():
      n = n + 1
      args = tests[test]
      expected_rtn = args[0]
      print "-----------------------------------------------------------------"
      print "-- Test %d: %s" % (n,test)
      print "--   ",args[1:]
      print "--    Expected return code: %s" % expected_rtn
      rtn = main(args[1:])
      print "-- Return code: %s" % rtn
      if ( rtn == expected_rtn ):
        print "-- Test %d: %s - SUCCESSFUL" % (n,test)
        print "-----------------------------------------------------------------"
      else:
        raise ConfigurationError("-- Test %d: %s - FAILED" % (n,test))
  except:
    raise
  print "**********## All %d tests passed ***************" % n

#--------------------
def validate_args(opts,expected_args):
  if ( len(opts) < expected_args ):
    raise UsageError("Insufficient number of arguments for option selected")
#---------------------
def show_line():
    x = traceback.extract_tb(sys.exc_info()[2])
    z = x[len(x)-1]
    return "%s line %s" % (z[2],z[1])

############################# Main Program ##############################

def main(opts=None):
  try:
      #--- process command line arguments ---
      ##print len(opts)
      validate_args(opts,2)
      opt = opts[1]
      if (opt == "-h") or (opt == "--help"):
        usage(opts[0]);return 1
      elif (opt == "--compare") or (opt == "-compare"):
        validate_args(opts,4)
        return compare_ini_files(opts[2],opts[3],False)
      elif (opt == "--compare-no-local") or (opt == "-compare-no-local"):
        validate_args(opts,4)
        return compare_ini_files(opts[2],opts[3],True)
      elif (opt == "--show-options") or (opt == "-show-options"):
        validate_args(opts,3)
        ini = Configuration(opts[2])
        print ini.section_options()
      elif (opt == "--show-values") or (opt == "-show-values"):
        validate_args(opts,3)
        ini = Configuration(opts[2])
        print ini
        validate_args(opts,3)
      elif (opt == "--validate") or (opt == "-validate"):
        ini = Configuration(opts[2])
        print "... configuration ini file has no syntax errors"
      elif (opt == "--test") or (opt == "-test"):
        run_unit_tests(opts[0])
      else:
        raise UsageError("Invalid command line option")
  except ConfigurationError, e:
    print;print "Configuration ERROR: %s" % e;return 1
  except UsageError, e:
    usage(opts[0])
    print "Usage ERROR: %s" % e;return 1
  except Exception, e:
    print;print "Exception ERROR: %s - %s" % (show_line(),e);return 1
  return 0

#--------------------------
if __name__ == '__main__':
  sys.exit(main(sys.argv))

<commit_msg>This fix is to resolve a problem with python's os.path.basename and os.path.dirname behavior as they do not work the same as the linux basename/dirname commands.  If the path specified contains a trailing '/', basename will be empty and dirname will give you the lowest level in the path rather than the parent as linux does. . This caused problems in some of the code which was attempting to validate the path and create missing directories. . It also caused a problem with the condor switchboard when privilege separation was in effect when the valid-dirs option value of the /etc/condor/privsep_config contained a trailing '/'.  The switchboard would throw an error saying it could not create the directory.  This occured when trying to create the factory. . To resolve this, the very basic method to retrieve a value for an ini option was changed to strip all trailing '/'s from all values returned.  This will work for the glideinWMS application as it stands right now.  If this changes and some other option does require or allow a trailing '/', this will have to change. . Whew! A lot of words for a small code change, but I thought it best that all are aware of this.<commit_after>#!/usr/bin/python

import common
#-----
import re
import os
import string
import sys
import time
import traceback
import ConfigParser
import StringIO

import inspect

#### Class Configuration ############################
class Configuration:

  def __init__ (self, fileName):
    self.cp       = None
    self.inifile = fileName

    #-- check for duplicate sections in config file --
    self.check_for_duplicate_sections()

    try:
      #-- check version of parser to use ---
      if sys.version_info[0] >= 2 and sys.version_info[1] >= 3:
        self.cp = ConfigParser.SafeConfigParser()
      else:
        self.cp = ConfigParser.ConfigParser()
      #-- read the ini file ---
      real_fp = open(self.inifile, 'r')
      string_fp = StringIO.StringIO(real_fp.read())
      self.cp.readfp(string_fp,self.inifile)
    except Exception, e:
      common.logerr("%s" % e)

    #-- additional check for syntax errors --
    self.syntax_check()

  #----------------
  def syntax_check(self):
    """ Checks for some syntax errors in ini config file. """
    for section in self.sections():
      for option in self.options(section):
        value = self.option_value(section,option)
        if "\n" in value:
          line = string.split(value,"\n")
          common.logerr("Section [%s]: this line starts with whitespace ( %s)\n       Please remove the leading space or comment (;) the line." % (section,line[1]))

  #----------------
  def __str__ (self):
    result = []
    result.append('<Configuration from %s>' % self.inifile)
    for section in self.sections():
      result.append('[%s]' % section)
      for option in self.options(section):
        value = self.option_value(section, option)
        result.append('    %-25s %s' % (option, value))
    return '\n'.join(result)

  #----------------
  def check_for_duplicate_sections(self):
    """ Check for duplicate sections in a config file ignoring commented ones.
        In addition, checks to verify there is no whitespace preceding or
        appending the section name in brackets as the ini parser does not
        validate for this.
    """
    if (self.inifile == "") or (self.inifile is None):
      common.logerr("System error: config file name is empty")
    try:
      fp = open(self.inifile, 'r')
    except:
      common.logerr("Problem reading ini file: %s" % sys.exc_info()[1])
    sections = {}    # used to keep track of sections
    duplicates = []  # used to identify duplicates
    for line in fp.readlines():
      newline = line.strip()
      if len(newline) == 0:
        continue
      if newline[0] != "[":
        continue
      #--- see if it is a section ---
      match = re.search('\[(.*)\]', newline)
      if match:
        section = match.group(1).lower().strip()
        if section in sections:
          duplicates.append(section)
          continue
        sections[section] = True
    if (len(duplicates) <> 0 ):
      common.logerr("Duplicate sections in %s - %s" % (self.inifile,duplicates))

  #----------------
  def validate_section(self,section,valid_option_list):
    if not self.has_section(section):
      common.logerr("Section (%s) does not exist in ini file (%s)" % (section,self.inifile))
    errors = []
    for option in valid_option_list:
      if self.has_option(section,option):
        continue
      errors.append(option)
    if len(errors) > 0:
      common.logerr("These options are not defined in the %s section of the ini file: %s" % (section,errors))

  #----------------
  def section_options(self):
    result = []
    for section in self.sections():
      result.append('[%s]' % section)
      for option in self.options(section):
        result.append('    %-25s' % (option))
    return '\n'.join(result)

  #----------------
  def filename(self):
    return self.inifile

  #----------------
  def sections(self):
    sections = list(self.cp.sections())
    sections.sort()
    return sections

  #----------------
  def options(self,section):
    options = self.cp.options(section)
    options.sort()
    return options

  #----------------
  def option_value(self,section,option):
    """ Due they way python os.path.basename/dirname work, we cannot let a
        pathname end in a '/' or we may see inconsistent results.  So we
        are stripping all option values of trailing '/'s.
    """
    value = self.cp.get(section,option)
    #-- cannot let paths end in a '/' --
    while len(value) > 0 and value[len(value)-1] == "/":
      value = value[0:len(value)-1].strip()
    return value

  #----------------
  def has_section(self, section):
    return self.cp.has_section(section)
  #----------------
  def has_option(self, section,option):
    return self.cp.has_option(section,option)
  #----------------
  def delete_section(self, section):
    self.cp.remove_section(section)
    return

#### Exceptions #####################################
class ConfigurationError(Exception):
  pass
class UsageError(Exception):
  pass

#####################################################
#---------------------
#---- Functions ------
#---------------------
def compare_ini_files (file_1, file_2, no_local_settings):
  try:
    print "Comparing ini files: %s / %s" % (file_1,file_2)
    ini_1 = Configuration(file_1)
    ini_2 = Configuration(file_2)
    rtn = 0
    #--- remove the Local Settings section conditionally ---
    if ( no_local_settings ):
      ini_1.delete_section("Local Settings")
      ini_2.delete_section("Local Settings")
    print "... Checking section information:"
    if ( ini_1.sections() == ini_2.sections() ):
      print "... sections are identical"
    else:
      print "... WARNING: section information differs"
      compare_sections(ini_1,ini_2)
      compare_sections(ini_2,ini_1)
      rtn = 1
    print
    print "... Checking section/object information:"
    if ( ini_1.section_options() ==  ini_2.section_options() ):
      print "... all section/objects are identical"
    else:
      print "... WARNING: section/object information differs"
      compare_options(ini_1,ini_2)
      compare_options(ini_2,ini_1)
      rtn = 1
  except:
    raise
  return rtn

#--------------------------------
def compare_sections(ini1,ini2):
  print """
... Sections     in %s
       NOT FOUND in %s""" % (ini1.filename(),ini2.filename())
  for section in ini1.sections():
    if ( ini2.has_section(section) ):
      continue
    else:
      print "    %s" % (section)

#--------------------------------
def compare_options(ini1,ini2):
  print """
... Section/objects in %s
          NOT FOUND in %s""" % (ini1.filename(),ini2.filename())
  for section in ini1.sections():
    for option in ini1.options(section):
      ## if (section == "SE CHANGEME"):
      ##   if (option == "enable"):
      ##     print section,option
      if ( ini2.has_option(section,option) == False):
        print "    %-20s/%s" % (section,option)

#--------------------------------
def usage(pgm):
  print
  print "Usage: " + pgm + " --compare file1 file2"
  print "       " + pgm + " --show-options file"
  print "       " + pgm + " --show-values file"
  print "       " + pgm + " --validate file"
  print "       " + pgm + " --help | -h "
  print """
   compare .......... Shows the differences between the
                      section/objects (not values) of the 2 ini files
                        returns 0 if identical
                        returns 1 if any differences
   show-options ...... Shows the section/objects for the ini file
   show-values ....... Shows the section/objects/values for the ini file
   validate .......... Verifies the ini file has no syntax errors

   Full path to the files must be specified unless this is executed
   in the directory in which they reside.

"""

#----------------------------------
def run_unit_tests(pgm):
  try:
    dir="./testdata/"
    tests = {
"no arguments"        : [1,pgm],
"not enough arguments": [1,pgm,"--validate"],
"not enough arguments": [1,pgm,"--compare"],
"invalid argument"    : [1,pgm,"--bad-arg",dir+"non-existent-file"],

##  validate ###
"validate: good ini"  : [0,pgm,"--validate",dir+"config-good.ini"],
"validate: bad ini"   : [1,pgm,"--validate",dir+"config-bad.ini"],
"validate: no ini"    : [1,pgm,"--validate",dir+"non-existent-file"],
"duplicates"          : [1,pgm,"--validate",dir+"config-w-dup-sections.ini"],

##  compare ###
"compare: no difference" : [0,pgm,"--compare",dir+"config-good.ini",dir+"config-good.ini"],
"compare: differences"   : [1,pgm,"--compare",dir+"config-good.ini",dir+"config-good-different.ini"],
"compare: no ini"        : [1,pgm,"--validate",dir+"non-existent-file"],
"duplicates"             : [1,pgm,"--compare",dir+"config-good.ini config-w-dup-sections.ini"],

##  show-options ###
"show-options"          : [0,pgm,"--show-options",dir+"config-good.ini"],
"show-options: bad ini" : [1,pgm,"--show-options",dir+"config-bad.ini"],
"show-options: no ini"  : [1,pgm,"--show-options",dir+"non-existent-file"],
"duplicates"            : [1,pgm,"--show-options",dir+"config-w-dup-sections.ini"],

##  show-values ###
"show-values"           : [0,pgm,"--show-values",dir+"config-good.ini"],
"show-values: bad ini"  : [1,pgm,"--show-values",dir+"config-bad.ini"],
"show-values: no ini"   : [1,pgm,"--show-values",dir+"non-existent-file"],
"duplicates"            : [1,pgm,"--show-values",dir+"config-w-dup-sections.ini"],
    }
    #---- run tests -----
    n=0
    for test in tests.keys():
      n = n + 1
      args = tests[test]
      expected_rtn = args[0]
      print "-----------------------------------------------------------------"
      print "-- Test %d: %s" % (n,test)
      print "--   ",args[1:]
      print "--    Expected return code: %s" % expected_rtn
      rtn = main(args[1:])
      print "-- Return code: %s" % rtn
      if ( rtn == expected_rtn ):
        print "-- Test %d: %s - SUCCESSFUL" % (n,test)
        print "-----------------------------------------------------------------"
      else:
        raise ConfigurationError("-- Test %d: %s - FAILED" % (n,test))
  except:
    raise
  print "**********## All %d tests passed ***************" % n

#--------------------
def validate_args(opts,expected_args):
  if ( len(opts) < expected_args ):
    raise UsageError("Insufficient number of arguments for option selected")
#---------------------
def show_line():
    x = traceback.extract_tb(sys.exc_info()[2])
    z = x[len(x)-1]
    return "%s line %s" % (z[2],z[1])

############################# Main Program ##############################

def main(opts=None):
  try:
      #--- process command line arguments ---
      ##print len(opts)
      validate_args(opts,2)
      opt = opts[1]
      if (opt == "-h") or (opt == "--help"):
        usage(opts[0]);return 1
      elif (opt == "--compare") or (opt == "-compare"):
        validate_args(opts,4)
        return compare_ini_files(opts[2],opts[3],False)
      elif (opt == "--compare-no-local") or (opt == "-compare-no-local"):
        validate_args(opts,4)
        return compare_ini_files(opts[2],opts[3],True)
      elif (opt == "--show-options") or (opt == "-show-options"):
        validate_args(opts,3)
        ini = Configuration(opts[2])
        print ini.section_options()
      elif (opt == "--show-values") or (opt == "-show-values"):
        validate_args(opts,3)
        ini = Configuration(opts[2])
        print ini
        validate_args(opts,3)
      elif (opt == "--validate") or (opt == "-validate"):
        ini = Configuration(opts[2])
        print "... configuration ini file has no syntax errors"
      elif (opt == "--test") or (opt == "-test"):
        run_unit_tests(opts[0])
      else:
        raise UsageError("Invalid command line option")
  except ConfigurationError, e:
    print;print "Configuration ERROR: %s" % e;return 1
  except UsageError, e:
    usage(opts[0])
    print "Usage ERROR: %s" % e;return 1
  except Exception, e:
    print;print "Exception ERROR: %s - %s" % (show_line(),e);return 1
  return 0

#--------------------------
if __name__ == '__main__':
  sys.exit(main(sys.argv))

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
##############################################################################
#
#    Author: Guewen Baconnier
#    Copyright 2013 Camptocamp SA
#
#    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/>.
#
##############################################################################

import logging
import os
import threading
import time
import traceback
import uuid
from datetime import datetime, timedelta
from StringIO import StringIO

import openerp
import openerp.modules.registry as registry_module
from .queue import JobsQueue
from ..session import ConnectorSessionHandler
from .job import (OpenERPJobStorage,
                  RETRY_INTERVAL,
                  ENQUEUED,
                  STARTED,
                  FAILED,
                  DONE)
from ..exception import (NoSuchJobError,
                         NotReadableJobError,
                         RetryableJobError,
                         FailedJobError,
                         NothingToDoJob)

_logger = logging.getLogger(__name__)

WAIT_CHECK_WORKER_ALIVE = 5  # seconds TODO change to 30 seconds
WAIT_WHEN_ONLY_AFTER_JOBS = 10  # seconds

class Worker(threading.Thread):
    """ Post and retrieve jobs from the queue, execute them"""

    queue_class = JobsQueue
    job_storage_class = OpenERPJobStorage

    def __init__(self, db_name, watcher):
        super(Worker, self).__init__()
        self.queue = self.queue_class()
        self.db_name = db_name
        self.uuid = unicode(uuid.uuid4())
        self.watcher = watcher

    def run_job(self, job):
        """ Execute a job """
        session_hdl = ConnectorSessionHandler(self.db_name,
                                          openerp.SUPERUSER_ID)
        try:
            job = self._load_job(job)
            if job is None:
                return

            # if the job has been manually set to DONE
            # before its execution, stop
            if job.state == DONE:
                return

            if job.eta and job.eta > datetime.now():
                # The queue is sorted by 'eta' date first
                # so if we dequeued a job expected to be run in
                # the future, we have no jobs to do right now!
                self.queue.enqueue(job)
                # Wait some time just to avoid to loop over
                # the same 'future' jobs
                _logger.debug('Wait %s seconds because the delayed '
                              'jobs have been reached',
                              WAIT_WHEN_ONLY_AFTER_JOBS)
                time.sleep(WAIT_WHEN_ONLY_AFTER_JOBS)
                return

            with session_hdl.session() as session:
                job.set_state(STARTED)
                self.job_storage_class(session).store(job)

            _logger.debug('%s started', job)
            with session_hdl.session() as session:
                job.perform(session)
            _logger.debug('%s done', job)

            with session_hdl.session() as session:
                job.set_state(DONE)
                self.job_storage_class(session).store(job)

        except NothingToDoJob:
            job.cancel()
            with session_hdl.session() as session:
                self.job_storage_class(session).store(job)

        except RetryableJobError:
            # delay the job later
            job.eta = timedelta(seconds=RETRY_INTERVAL)
            with session_hdl.session() as session:
                self.job_storage_class(session).store(job)

        except (FailedJobError, Exception):
            buff = StringIO()
            traceback.print_exc(file=buff)
            _logger.error(buff.getvalue())

            job.set_state(FAILED, exc_info=buff.getvalue())
            with session_hdl.session() as session:
                self.job_storage_class(session).store(job)
            raise

    def _load_job(self, job):
        """ Reload a job from the backend """
        session_hdl = ConnectorSessionHandler(self.db_name,
                                              openerp.SUPERUSER_ID)
        with session_hdl.session() as session:
            try:
                job = self.job_storage_class(session).load(job.uuid)
            except NoSuchJobError:
                # just skip it
                job = None
            except NotReadableJobError:
                _logger.exception('Could not read job: %s', job)
                raise
        return job

    def run(self):
        """ Worker's main loop

        Check if it still exists in the ``watcher``. When it does no
        longer exist, it break the loop so the thread stops properly.

        Wait for jobs and execute them sequentially.
        """
        while 1:
            # check if the worker has to exit
            if self.watcher.worker_lost(self):
                break
            job = self.queue.dequeue()
            try:
                self.run_job(job)
            except:
                continue

    def enqueue_job_uuid(self, job_uuid):
        """ Enqueue a job:

        It will be executed by the worker as soon as possible (according
        to the job's priority
        """
        session_hdl = ConnectorSessionHandler(self.db_name,
                                              openerp.SUPERUSER_ID)
        with session_hdl.session() as session:
            try:
                job = self.job_storage_class(session).load(job_uuid)
            except NoSuchJobError:
                # just skip it
                return
            except NotReadableJobError:
                _logger.exception('Could not read job: %s', job_uuid)
                raise
            job.set_state(ENQUEUED)
            self.job_storage_class(session).store(job)
        # the change of state should be commited before
        # the enqueue otherwise we may have concurrent updates
        self.queue.enqueue(job)
        _logger.debug('%s enqueued in %s', job, self)

class WorkerWatcher(threading.Thread):
    """ Keep a sight on the workers and signal their aliveness.

    A `WorkerWatcher` is shared between databases, so only 1 instance is
    necessary to check the aliveness of the workers for every database.
    """

    def __init__(self):
        super(WorkerWatcher, self).__init__()
        self.workers = {}

    def new(self, db_name):
        """ Create a new worker for the database """
        if db_name in self.workers:
            raise Exception('Database %s already has a worker (%s)' %
                            (db_name, self.workers[db_name].uuid))
        worker = Worker(db_name, self)
        self.workers[db_name] = worker
        worker.daemon = True
        worker.start()

    def delete(self, db_name):
        """ Delete worker for the database """
        if db_name in self.workers:
            del self.workers[db_name]

    def worker_lost(self, worker):
        """ Indicate if a worker is no longer referenced by the watcher.

        Used by the worker threads to know they have to exit.
        """
        return worker not in self.workers.itervalues()

    @staticmethod
    def available_registries():
        """ Yield the registries which are available.

        Available means that they can be used by a `Worker`.

        :return: database name, registry
        :rtype: tuple
        """
        registries = registry_module.RegistryManager.registries
        for db_name, registry in registries.iteritems():
            if not 'connectors.installed' in registry.models:
                continue
            if not registry.ready:
                continue
            yield db_name, registry

    def _update_workers(self):
        """ Refresh the list of workers according to the available
        databases and registries.

        A new database can be available, so we need to create a new
        `Worker` or a database could have been dropped, so we have to
        discard the Worker.
        """
        for db_name, _registry in self.available_registries():
            if db_name not in self.workers:
                self.new(db_name)

        # XXX not necessary if we keep the monkey patch of
        # RegistryManager.delete
        all_db = registry_module.RegistryManager.registries.keys()
        for removed_db in set(self.workers) ^ set(all_db):
            self.delete(removed_db)

    def run(self):
        """ `WorkerWatcher`'s main loop """
        while 1:
            self._update_workers()
            for db_name, worker in self.workers.items():
                self.check_alive(db_name, worker)
            time.sleep(WAIT_CHECK_WORKER_ALIVE)

    def check_alive(self, db_name, worker):
        """ Check if the the worker is still alive and notify
        its aliveness.
        Check if the other workers are still alive, if they are
        dead, remove them from the worker's pool.
        """
        session_hdl = ConnectorSessionHandler(db_name,
                                              openerp.SUPERUSER_ID)
        with session_hdl.session() as session:
            if worker.is_alive():
                self._notify_alive(session, worker)
                session.commit()
            self._purge_dead_workers(session)
            session.commit()

    def _notify_alive(self, session, worker):
        _logger.debug('Worker %s is alive on process %s',
                      worker.uuid, os.getpid())
        dbworker_obj = session.pool.get('queue.worker')
        dbworker_obj._notify_alive(session.cr,
                                   session.uid,
                                   worker,
                                   context=session.context)

    def _purge_dead_workers(self, session):
        dbworker_obj = session.pool.get('queue.worker')
        dbworker_obj._purge_dead_workers(session.cr,
                                         session.uid,
                                         context=session.context)

watcher = WorkerWatcher()

registry_delete_original = registry_module.RegistryManager.delete
def delete(cls, db_name):
    """Delete the registry linked to a given database.  """
    watcher.delete(db_name)
    return registry_delete_original(db_name)
registry_module.RegistryManager.delete = classmethod(delete)

def start_service():
    """ Start the watcher """
    watcher.daemon = True
    watcher.start()

start_service()
<commit_msg>check the workers aliveness every 30 seconds<commit_after># -*- coding: utf-8 -*-
##############################################################################
#
#    Author: Guewen Baconnier
#    Copyright 2013 Camptocamp SA
#
#    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/>.
#
##############################################################################

import logging
import os
import threading
import time
import traceback
import uuid
from datetime import datetime, timedelta
from StringIO import StringIO

import openerp
import openerp.modules.registry as registry_module
from .queue import JobsQueue
from ..session import ConnectorSessionHandler
from .job import (OpenERPJobStorage,
                  RETRY_INTERVAL,
                  ENQUEUED,
                  STARTED,
                  FAILED,
                  DONE)
from ..exception import (NoSuchJobError,
                         NotReadableJobError,
                         RetryableJobError,
                         FailedJobError,
                         NothingToDoJob)

_logger = logging.getLogger(__name__)

WAIT_CHECK_WORKER_ALIVE = 30  # seconds
WAIT_WHEN_ONLY_AFTER_JOBS = 10  # seconds

class Worker(threading.Thread):
    """ Post and retrieve jobs from the queue, execute them"""

    queue_class = JobsQueue
    job_storage_class = OpenERPJobStorage

    def __init__(self, db_name, watcher):
        super(Worker, self).__init__()
        self.queue = self.queue_class()
        self.db_name = db_name
        self.uuid = unicode(uuid.uuid4())
        self.watcher = watcher

    def run_job(self, job):
        """ Execute a job """
        session_hdl = ConnectorSessionHandler(self.db_name,
                                          openerp.SUPERUSER_ID)
        try:
            job = self._load_job(job)
            if job is None:
                return

            # if the job has been manually set to DONE
            # before its execution, stop
            if job.state == DONE:
                return

            if job.eta and job.eta > datetime.now():
                # The queue is sorted by 'eta' date first
                # so if we dequeued a job expected to be run in
                # the future, we have no jobs to do right now!
                self.queue.enqueue(job)
                # Wait some time just to avoid to loop over
                # the same 'future' jobs
                _logger.debug('Wait %s seconds because the delayed '
                              'jobs have been reached',
                              WAIT_WHEN_ONLY_AFTER_JOBS)
                time.sleep(WAIT_WHEN_ONLY_AFTER_JOBS)
                return

            with session_hdl.session() as session:
                job.set_state(STARTED)
                self.job_storage_class(session).store(job)

            _logger.debug('%s started', job)
            with session_hdl.session() as session:
                job.perform(session)
            _logger.debug('%s done', job)

            with session_hdl.session() as session:
                job.set_state(DONE)
                self.job_storage_class(session).store(job)

        except NothingToDoJob:
            job.cancel()
            with session_hdl.session() as session:
                self.job_storage_class(session).store(job)

        except RetryableJobError:
            # delay the job later
            job.eta = timedelta(seconds=RETRY_INTERVAL)
            with session_hdl.session() as session:
                self.job_storage_class(session).store(job)

        except (FailedJobError, Exception):
            buff = StringIO()
            traceback.print_exc(file=buff)
            _logger.error(buff.getvalue())

            job.set_state(FAILED, exc_info=buff.getvalue())
            with session_hdl.session() as session:
                self.job_storage_class(session).store(job)
            raise

    def _load_job(self, job):
        """ Reload a job from the backend """
        session_hdl = ConnectorSessionHandler(self.db_name,
                                              openerp.SUPERUSER_ID)
        with session_hdl.session() as session:
            try:
                job = self.job_storage_class(session).load(job.uuid)
            except NoSuchJobError:
                # just skip it
                job = None
            except NotReadableJobError:
                _logger.exception('Could not read job: %s', job)
                raise
        return job

    def run(self):
        """ Worker's main loop

        Check if it still exists in the ``watcher``. When it does no
        longer exist, it break the loop so the thread stops properly.

        Wait for jobs and execute them sequentially.
        """
        while 1:
            # check if the worker has to exit
            if self.watcher.worker_lost(self):
                break
            job = self.queue.dequeue()
            try:
                self.run_job(job)
            except:
                continue

    def enqueue_job_uuid(self, job_uuid):
        """ Enqueue a job:

        It will be executed by the worker as soon as possible (according
        to the job's priority
        """
        session_hdl = ConnectorSessionHandler(self.db_name,
                                              openerp.SUPERUSER_ID)
        with session_hdl.session() as session:
            try:
                job = self.job_storage_class(session).load(job_uuid)
            except NoSuchJobError:
                # just skip it
                return
            except NotReadableJobError:
                _logger.exception('Could not read job: %s', job_uuid)
                raise
            job.set_state(ENQUEUED)
            self.job_storage_class(session).store(job)
        # the change of state should be commited before
        # the enqueue otherwise we may have concurrent updates
        self.queue.enqueue(job)
        _logger.debug('%s enqueued in %s', job, self)

class WorkerWatcher(threading.Thread):
    """ Keep a sight on the workers and signal their aliveness.

    A `WorkerWatcher` is shared between databases, so only 1 instance is
    necessary to check the aliveness of the workers for every database.
    """

    def __init__(self):
        super(WorkerWatcher, self).__init__()
        self.workers = {}

    def new(self, db_name):
        """ Create a new worker for the database """
        if db_name in self.workers:
            raise Exception('Database %s already has a worker (%s)' %
                            (db_name, self.workers[db_name].uuid))
        worker = Worker(db_name, self)
        self.workers[db_name] = worker
        worker.daemon = True
        worker.start()

    def delete(self, db_name):
        """ Delete worker for the database """
        if db_name in self.workers:
            del self.workers[db_name]

    def worker_lost(self, worker):
        """ Indicate if a worker is no longer referenced by the watcher.

        Used by the worker threads to know they have to exit.
        """
        return worker not in self.workers.itervalues()

    @staticmethod
    def available_registries():
        """ Yield the registries which are available.

        Available means that they can be used by a `Worker`.

        :return: database name, registry
        :rtype: tuple
        """
        registries = registry_module.RegistryManager.registries
        for db_name, registry in registries.iteritems():
            if not 'connectors.installed' in registry.models:
                continue
            if not registry.ready:
                continue
            yield db_name, registry

    def _update_workers(self):
        """ Refresh the list of workers according to the available
        databases and registries.

        A new database can be available, so we need to create a new
        `Worker` or a database could have been dropped, so we have to
        discard the Worker.
        """
        for db_name, _registry in self.available_registries():
            if db_name not in self.workers:
                self.new(db_name)

        # XXX not necessary if we keep the monkey patch of
        # RegistryManager.delete
        all_db = registry_module.RegistryManager.registries.keys()
        for removed_db in set(self.workers) ^ set(all_db):
            self.delete(removed_db)

    def run(self):
        """ `WorkerWatcher`'s main loop """
        while 1:
            self._update_workers()
            for db_name, worker in self.workers.items():
                self.check_alive(db_name, worker)
            time.sleep(WAIT_CHECK_WORKER_ALIVE)

    def check_alive(self, db_name, worker):
        """ Check if the the worker is still alive and notify
        its aliveness.
        Check if the other workers are still alive, if they are
        dead, remove them from the worker's pool.
        """
        session_hdl = ConnectorSessionHandler(db_name,
                                              openerp.SUPERUSER_ID)
        with session_hdl.session() as session:
            if worker.is_alive():
                self._notify_alive(session, worker)
                session.commit()
            self._purge_dead_workers(session)
            session.commit()

    def _notify_alive(self, session, worker):
        _logger.debug('Worker %s is alive on process %s',
                      worker.uuid, os.getpid())
        dbworker_obj = session.pool.get('queue.worker')
        dbworker_obj._notify_alive(session.cr,
                                   session.uid,
                                   worker,
                                   context=session.context)

    def _purge_dead_workers(self, session):
        dbworker_obj = session.pool.get('queue.worker')
        dbworker_obj._purge_dead_workers(session.cr,
                                         session.uid,
                                         context=session.context)

watcher = WorkerWatcher()

registry_delete_original = registry_module.RegistryManager.delete
def delete(cls, db_name):
    """Delete the registry linked to a given database.  """
    watcher.delete(db_name)
    return registry_delete_original(db_name)
registry_module.RegistryManager.delete = classmethod(delete)

def start_service():
    """ Start the watcher """
    watcher.daemon = True
    watcher.start()

start_service()
<|endoftext|>