write_excel.py

"""
Contains utility class for writing to xlsx
"""

from openpyxl import Workbook, load_workbook
from openpyxl.utils import get_column_letter
from openpyxl.styles import Font
from openpyxl.styles.fills import PatternFill
from openpyxl.chart import PieChart, Reference

class ExcelWriter:
    """
    Utility class which writes `TestEntry` objects to an excel file.
    Creates file if it does not exist.
    """
    PASS_COL = "00FF00"
    FAIL_COL = "FF0000"
    def __init__(self, output_file):
        self.output_file = output_file
        self.work_book = self.get_workbook()
        self.work_sheet = self.work_book.active

    def get_entry_with_id(self, test_id):
        """
        When inserting an entry, it might be a test being re-run. In that case,
        the row containing that test needs to be updated, as opposed to appending
        the entry to the end of the file.
        This method finds that row, returning the row number,
        or -1 in the case where that test is not in the report yet.
        """
        for cell in self.work_sheet["A"]:
            if cell.value is None:
                return -1
            if cell.value == test_id:
                print("MATCH with id {} at row {}".format(test_id, cell.row))
                return cell.row
        return -1

    def get_last_row(self):
        """
        return index of first empty row
        """
        for cell in self.work_sheet["A"]:
            if cell.value is None:
                return cell.row
        return cell.row + 1

    def write_test_entry(self, test_entry):
        """
        Write a test entry to the work_sheet
        """
        existing_entry_row = self.get_entry_with_id(test_entry.test_id)
        last_row = self.get_last_row()
        # Use the above two values to pick a row
        entry_row = existing_entry_row if existing_entry_row != -1 else last_row

        # Pick a cell color based on test outcome
        cell_col = ExcelWriter.PASS_COL if test_entry.result == "PASS" else ExcelWriter.FAIL_COL

        # Test entry as a list
        entry_vals = test_entry.as_list()

        for col, cell_vale in zip(self.work_sheet.iter_cols(min_row=entry_row,
                                                            max_col=len(entry_vals),
                                                            max_row=entry_row), entry_vals):
            for cell in col:
                cell.value = cell_vale
                cell.fill = PatternFill("solid", fgColor=cell_col)# cellCol

    def get_workbook(self):
        """
        Returns workbook at output_file/.
        Creates and initialises it (adds headers) if it doesn't exist yet.
        """
        try:
            work_book = load_workbook(filename=self.output_file)
            return work_book
        except FileNotFoundError:
            work_book = Workbook()
            init_workbook(work_book)
            return work_book
    
    def get_chart_work_sheet(self):
        """
        Returns work sheet where charts will be drawn to.
        Creates if does not exist
        """
        try:
            chart_work_sheet = self.work_book["charts"]
            return chart_work_sheet
        except KeyError:
            return self.work_book.create_sheet("charts")
    
    def write_pie_chart_data(self):
        """
        Writes pie chart data (the charts sheet) required to construct the pie chart later
        """
        chart_work_sheet = self.get_chart_work_sheet()

        # Get data from main work sheet
        last_row = self.get_last_row()
        data = Reference(self.work_sheet, min_col=3, max_col=3, min_row=2, max_row=last_row)

        # Construct pie chart
        pie = PieChart()
        pie.add_data(data, titles_from_data=True)
        #pie.set_categories(labels)
        pie.title = "PASS/FAIL Distribution"

        chart_work_sheet.add_chart(pie, "A1")
    
    def write_pie_chart(self):
        """
        Writes a pie chart showing PASS/FAIL stats.
        Writes the chart to a separate work sheet.
        """
        chart_work_sheet = self.get_chart_work_sheet()

        # Get data from main work sheet
        last_row = self.get_last_row()
        data = Reference(self.work_sheet, min_col=3, max_col=3, min_row=2, max_row=last_row)

        # Construct pie chart
        pie = PieChart()
        pie.add_data(data, titles_from_data=True)
        #pie.set_categories(labels)
        pie.title = "PASS/FAIL Distribution"

        chart_work_sheet.add_chart(pie, "A1")

    def save(self):
        """
        Save workbook to disk.
        """
        self.work_book.save(filename=self.output_file)

def init_workbook(work_book):
    """
    Writes column headers to ws.
    """
    headers = [("Test ID", 10), ("Test name", 80), ("Result", 6), ("Error Message", 100),
               ("NFV API", 25), ("Robot Test File", 25)]
    header_font = Font(bold=True)
    work_sheet = work_book.active
    for col, header in zip(work_sheet.iter_cols(min_row=1, max_col=len(headers),
                                                max_row=1), headers):
        for cell in col:
            header_name = header[0]
            colsize = header[1]
            work_sheet.column_dimensions[get_column_letter(cell.column)].width = colsize
            cell.value = header_name
            cell.font = header_font