wanzb_v3_api/src/utils/plottool.py

# coding: utf-8
import re
import StringIO
import uuid
import logging

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 plot_records_as_table_and_convert_to_png(
    (nrows, ncols),
    data,
    col_headers=None,
    title=None):
    """
    visual records
    """
    fig_width = ncols
    offset = 0
    if col_headers:
        offset += 1
    if title:
        offset += 1
    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 * 2, edgecolor='red', facecolor='red', loc='right')
        table[0, 6].get_text().set_text(title)
        table[0, 6].set_fontsize(14)
        table[0, 6].get_text().set_color('yellow')

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

    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(10)
            if not table_data[i][5].startswith('-'):
                table[i + offset, j].get_text().set_color('red')

    ax.add_table(table)
    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 = []
    start = 0
    step = 100
    while True:
        part = data[start: start + step]
        part_nrows = len(part)

        if start == 0:
            png_bytes = plot_records_as_table_and_convert_to_png(
                (part_nrows, ncols),
                part,
                col_headers,
                title)
        else:
            png_bytes = plot_records_as_table_and_convert_to_png(
                (part_nrows, ncols),
                part)
        filename = u'{}-第{}-{}名.png'.format(title, start + 1, start + part_nrows)
        logging.error(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

    return ret