""" This module contains classes for creating a report spreadsheet with data. ReportBuilder class is used to create it. The specs of sheets are defined in the BaseSheet class and its subclasses. """ from typing import Callable, Collection, Generator, Iterable from openpyxl import Workbook from openpyxl.cell.cell import Cell from openpyxl.formatting.rule import FormulaRule from openpyxl.styles import Font, PatternFill from openpyxl.utils import get_column_letter from openpyxl.worksheet.datavalidation import DataValidation from openpyxl.worksheet.worksheet import Worksheet from .....typings import ColumnIndex, Numeric from .....utils import xlsx from .....utils.performance import no_garbage_collection from .sheets import BaseSheet, Column class ReportBuilder(Workbook): """Class which serves for creating a report spreadsheet. It extends the openpyxl Workbook class and adds a method for adding a new sheet with data. If no sheets are added, the workbook will be empty. """ _default_font_settings: dict[str, str | int] = {"name": "Arial", "size": 10} _default_font_header: Font = Font(**_default_font_settings, bold=True) _default_font_rows: Font = Font(**_default_font_settings) _freeze_panes_at_cell: str = "A2" def __init__(self): """Initialize a new report workbook and remove the default active sheet, so that it starts empty. """ super().__init__() self.remove(self.active) def add_sheet(self, sheet: BaseSheet) -> None: """Adds a new sheet with data to the workbook. Args: sheet: Sheet instance to include in the report. """ ws = self.create_sheet(sheet.title) sheet_columns = sheet.columns dropdown_column_indices = self._get_indices_of_cols_with_predicate( sheet_columns, lambda column: column.is_dropdown ) dropdowns = _Dropdowns(ws) with no_garbage_collection(): ws.append([column.label for column in sheet_columns]) self._apply_style_to_header(ws, sheet_columns) for row in sheet.rows: row_data = [] for column in sheet_columns: value = self._get_column_value(column.key, row) if column.handler: value = column.handler(value) row_data.append(value) ws.append(row_data) for col_idx in dropdown_column_indices: has_default_option = sheet_columns[col_idx].dropdown_default dropdowns.add(col_idx + 1, default=has_default_option) # 1-based index self._apply_style_to_rows(ws) self.add_conditional_highlighting(ws, sheet_columns) self.handle_hyperlinks(ws, sheet_columns) # Auto-size columns ONLY once all rows have been added (it takes # into account the width of the cells with the longest content). for column in ws.columns: xlsx.column_autosize(ws, column) ws.freeze_panes = self._freeze_panes_at_cell @staticmethod def _get_column_value( key: str | Iterable[str], row: dict[str, str | None], ) -> str: """Returns the value from the row for the specified column key. If the column key is a string, it is used as a key to get the value from the row. If it is an iterable, the first truthy value from the row is used. """ if isinstance(key, str): value = row.get(key) elif isinstance(key, Iterable): first_truthy_value = next((row.get(k) for k in key if row.get(k))) value = first_truthy_value else: raise ValueError("Column key must be a string or an iterable.") return value or "" def _apply_style_to_rows(self, ws: Worksheet) -> None: """Applies style to all rows excluding header. Args: ws: Worksheet. """ start_row = 2 end_row = ws.max_row self._apply_height(ws, start_row, end_row, height=15) # Get all cells in the WS excluding header range_cells = ws.iter_rows( min_row=start_row, max_row=end_row, min_col=1, max_col=ws.max_column ) self._apply_font((cell for row in range_cells for cell in row)) @staticmethod def _apply_height( ws: Worksheet, start_row: Numeric, end_row: Numeric, *, height: int, ) -> None: """Applies the specified height to the rows in the specified range. Args: ws: Worksheet. start_row: Start row index (1-based), inclusive. end_row: End row index (1-based), inclusive. height: Height to apply to the rows in the range, in points. """ row_range = range(start_row, end_row + 1) for row_idx in row_range: ws.row_dimensions[row_idx].height = height def _apply_font(self, cells: Collection[Cell]) -> None: font = self._default_font_rows for cell in cells: cell.font = font def to_bytes(self) -> bytes: """Converts the workbook to a byte stream. Returns: bytes: Byte stream of the workbook. """ return xlsx.workbook_to_xlsx_bytes(self) @staticmethod def add_conditional_highlighting( ws: Worksheet, sheet_columns: Collection[Column] ) -> None: """Adds conditional highlighting to the worksheet columns which have the highlight_if attribute set. """ for idx, column in enumerate(sheet_columns): col_idx = idx + 1 if column.highlight_if: max_row = ws.max_row formulas = [ column.highlight_if(row_idx) for row_idx in range(2, max_row + 1) ] rule = FormulaRule( formula=formulas, fill=PatternFill( start_color=column.color_highlight, end_color=column.color_highlight, fill_type="solid", ), ) col_letter = get_column_letter(col_idx) col_range = f"{col_letter}2:{col_letter}{max_row}" # No header ws.conditional_formatting.add(col_range, rule) @staticmethod def handle_hyperlinks(ws: Worksheet, sheet_columns: Collection[Column]) -> None: """Handles hyperlinks in the specified columns. All columns with hyperlink attribute set to True will be scanned for hyperlinks. Thus, if a cell contains a URL, it will be converted to a hyperlink. Args: ws: Worksheet. sheet_columns: List of columns in the sheet. """ for col_idx, column in enumerate(sheet_columns, start=1): if not column.is_hyperlink: continue col_letter = get_column_letter(col_idx) for row_idx, cell in enumerate(ws[f"{col_letter}"], start=1): if value := cell.value: if value.startswith("http"): cell.hyperlink = value cell.style = "Hyperlink" @staticmethod def _get_indices_of_cols_with_predicate( sheet_columns: Collection[Column], predicate: Callable[[Column], bool] ) -> tuple[ColumnIndex, ...]: """Returns 0-based indices of columns which satisfy the predicate. Args: sheet_columns: the column objects. predicate: the predicate to check if the column index should be included. """ matched_column_indices = tuple( col_idx for col_idx, column in enumerate(sheet_columns) if predicate(column) ) return matched_column_indices def _apply_style_to_header( self, ws: Worksheet, sheet_columns: Collection[Column] ) -> None: """Applies style to the header row. Args: ws: Worksheet. sheet_columns: List of columns in the sheet. """ self._apply_height(ws, start_row=1, end_row=1, height=33.75) font = self._default_font_header for col_idx, column in enumerate(sheet_columns, start=1): cell = ws.cell(row=1, column=col_idx) fill_color = column.color_label_bkg cell.fill = PatternFill( start_color=fill_color, end_color=fill_color, fill_type="solid", ) cell.font = font class _Dropdowns: """Class for managing dropdowns in a worksheet.""" def __init__(self, ws: Worksheet, separator: str = ";") -> None: """Initializes the dropdowns manager. Args: ws: Worksheet. separator: Separator used to split dropdown options. The cell value will be split by this separator to get the dropdown options. If there are less than 2 options, the dropdown will not be added and the cell value will remain unchanged. """ self.ws = ws self.separator = separator def add(self, col_idx: ColumnIndex, default: bool = True) -> None: """Adds dropdowns to the specified column. The column will be scanned for cells with a value that contains the specified separator. If found, the string will be split and the first option will be set as the default value for the dropdown. If there is only one option, the dropdown will not be added. Args: col_idx: Column index. Must be 1-based. default: Whether to set the first option as the default value. If False, the dropdown will have no default value and will start blank. """ if col_idx < 1: raise ValueError("Column index must be 1-based.") cached_options = {} cells_ex_header = self._get_cells_ex_header(col_idx) for cell in cells_ex_header: if cached_option := cached_options.get(cell.value): options = cached_option else: options = self._parse_options(cell.value or "") cached_options[cell.value] = options if len(options) < 2: continue self._add_dropdown(cell, options, default) def _get_cells_ex_header(self, col_idx: ColumnIndex) -> Generator[Cell, None, None]: """Returns cells in the specified column, excluding the header.""" col_rows = self.ws.iter_rows(min_col=col_idx, max_col=col_idx, min_row=2) col_cells = (cells[0] for cells in col_rows) return col_cells def _parse_options( self, string: str, ) -> list[str]: """Parses the options from the string, using the separator. Empty options will be discarded. """ options = (option.strip() for option in string.split(self.separator)) non_empty_options = [option for option in options if option] return non_empty_options def _add_dropdown(self, cell: Cell, options: list[str], default: bool) -> None: """Adds a dropdown to the specified cell. Args: cell: Cell to add the dropdown to. options: List of dropdown options. The first option will be set as the default value. default: Whether to set the first option as the default value. If False, the dropdown will have no default value and will start blank. """ dropdown = DataValidation( type="list", formula1=f'"{",".join(options)}"', allowBlank=False ) self.ws.add_data_validation(dropdown) dropdown.add(cell) cell.value = options[0] if default else None