wanzb_v3_api/src/utils/plottool.py

# coding: utf-8
import re
import StringIO
import uuid
import logging
import zipfile,os
import tempfile

import matplotlib as mpl
mpl.use('Agg')
mpl.rcParams['font.family'] = ['SimHei', 'sans-serif']

import matplotlib.pyplot as plt
import matplotlib.table as tb
#from utils.upload_to_oss import hnoss

def get_zipfile(files):                                                             
    """                                                                            
    """                                                                             
    path = os.path.split(files[0])[0]                                               
    temp = tempfile.TemporaryFile()                                                 
    archive = zipfile.ZipFile(temp, 'w', zipfile.ZIP_DEFLATED)
    for f in files:
        fname = os.path.split(f)[-1]
        archive.write(f,fname)
        os.remove(f)
    archive.close()
    temp.seek(0)
    return temp


def plot_records_as_table_and_convert_to_png(
    (nrows, ncols),
    data,
    col_headers=None,
    title=None,localfile=None):
    """
    visual records
    """
    fig_width = ncols
    h_offset = t_offset = 0
    if col_headers:
        h_offset = 1
    if title:
        t_offset = 1
    offset = h_offset + t_offset
    fig_height = offset * 2

    table_data = []
    row_heights_mul = []
    for i in data:
        t = [j for j in i]
        t[-1] = re.sub(',', '\n', t[-1])
        t[-2] = re.sub(',', '\n', t[-2])
        nlines = max(t[-1].count('\n'), t[-2].count('\n'))
        row_heights_mul.append(nlines + 2)
        fig_height += (nlines + 2)
        table_data.append(t)

    fig = plt.figure(figsize=(fig_width, float(fig_height)/5))
    ax = fig.add_axes([0, 0, 1, 1])
    ax.axis('off')

    table = tb.Table(ax, bbox=[0, 0, 1, 1])
    table.auto_set_font_size(False)
    widths = [0.05, 0.15, 0.1, 0.1, 0.1, 0.1, 0.2, 0.2]
    height = table._approx_text_height()

    if title:
        for i in range(ncols):
            table.add_cell(0, i, width=widths[i], height=height * 4, edgecolor='red', facecolor='red', loc='right')
        table[0, 6].get_text().set_text(title)
        table[0, 6].set_fontsize(25)
        table[0, 6].set_text_props(weight='bold')
        table[0, 6].get_text().set_color('yellow')

    if col_headers:
        for i in range(ncols):
            table.add_cell(t_offset, i, width=widths[i], height=height * 2, facecolor='yellow', text=col_headers[i], loc='center')
            table[t_offset, i].set_fontsize(12)
            table[t_offset, i].set_edgecolor('grey')

    for i in range(nrows):
        for j in range(ncols):
            table.add_cell(i + offset, j, width=widths[j], height=height*row_heights_mul[i], text=table_data[i][j], loc='center')
            table[i + offset, j].set_fontsize(14)
            table[i + offset, j].set_edgecolor('grey')
            if not table_data[i][5].startswith('-'):
                table[i + offset, j].get_text().set_color('red')

    ax.add_table(table)
    if localfile:
        plt.savefig(localfile, format='png')
        return localfile
    else:
        buf = StringIO.StringIO()
        plt.savefig(buf, format='png')
        return buf.getvalue()


def plot_records(data, col_headers, title, prefix):
    nrows = len(data)
    if nrows == 0:
        return
    ncols = len(col_headers) if col_headers else len(data[0])

    ret = []
    files = []
    start = 0
    step = 120
    while True:
    	if start==0 and ('百万组' in title or '十万组' in title or '菜鸟组' in title ):
    		step = 10
    	else:
    		step = 120
    		
        part = data[start: start + step]
        part_nrows = len(part)
        filename = u'/tmp/{}-{}.png'.format(title.split('-')[1], start)

        if start == 0:
            png_bytes = plot_records_as_table_and_convert_to_png(
                (part_nrows, ncols),
                part,
                col_headers,
                title,filename)
        else:
            png_bytes = plot_records_as_table_and_convert_to_png(
                (part_nrows, ncols),
                part,None,None,filename)
        logging.error(filename)
        files.append(filename)
        #image_url = hnoss.upload_from_str(
        #    png_bytes,
        #    '{}_{}.png'.format(prefix, uuid.uuid4().hex),
        #    False)
        #ret.append({'filename': filename, 'image_url': image_url})

        start += step
        if start >= nrows:
            break
    zpfile = get_zipfile(files)
    return zpfile