pydantic_xlsx.model

"""
Provides the `XlsxModel` to the library user.
"""

from pathlib import Path
from typing import Any, List, Optional, Type, Union

from openpyxl import load_workbook, Workbook
from pydantic import BaseModel

from .composition import CompositionFactory
from .config import XlsxConfig


class XlsxModel(BaseModel):
    """
    Extends pydantic with the ability to in- and export data from/to xlsx
    files. The class also allows the usage of collection by using the __root__
    parameter.
    """

    Config = XlsxConfig

    def __init__(self, **data: Any):
        super().__init__(**data)

    @classmethod
    def from_file(
            cls: Type["XlsxModel"],
            path: Union[str, Path]
    ) -> "XlsxModel":
        """Loads data from a Excel file."""
        return cls.from_workbook(load_workbook(path))

    @classmethod
    def from_workbook(
            cls: Type["XlsxModel"],
            wb: Workbook
    ) -> "XlsxModel":
        """Loads data from a Excel file."""

        return CompositionFactory.from_model(cls).workbook_to_model(cls, wb)

    def workbook(self) -> Workbook:
        """Returns a openpyxl Workbook."""
        wb = Workbook()
        sheets = CompositionFactory.from_model(self).sheets_from_model(self)

        first_sheet = True
        for sheet in sheets:
            if first_sheet:
                ws = wb.active
                first_sheet = False
            else:
                ws = wb.create_sheet()
            sheet.populate_worksheet(ws)
        return wb

    def to_file(self, path: Union[str, Path]):
        """Saves the model to a Excel file."""
        wb = self.workbook()
        wb.save(path)

    @classmethod
    def _property_keys(cls) -> List[str]:
        """Returns the names of the properties used in the excel header."""
        return [field.alias for field in cls.__fields__.values()]

    @classmethod
    def _ignore_key(cls, key: Optional[str]) -> bool:
        """
        Takes a property name (key) and checks if ignore_additional_rows is
        enabled the property is described by the model. This method is used by
        the from_workbook methods to determine if a cell should be added to the
        import dict.
        """
        ignore = getattr(cls.__config__, "ignore_additional_rows", None)
        if ignore is not None and not ignore:
            return False
        if key is None:
            return True
        return key not in cls._property_keys()

class XlsxModel(BaseModel):
    """
    Extends pydantic with the ability to in- and export data from/to xlsx
    files. The class also allows the usage of collection by using the __root__
    parameter.
    """

    Config = XlsxConfig

    def __init__(self, **data: Any):
        super().__init__(**data)

    @classmethod
    def from_file(
            cls: Type["XlsxModel"],
            path: Union[str, Path]
    ) -> "XlsxModel":
        """Loads data from a Excel file."""
        return cls.from_workbook(load_workbook(path))

    @classmethod
    def from_workbook(
            cls: Type["XlsxModel"],
            wb: Workbook
    ) -> "XlsxModel":
        """Loads data from a Excel file."""

        return CompositionFactory.from_model(cls).workbook_to_model(cls, wb)

    def workbook(self) -> Workbook:
        """Returns a openpyxl Workbook."""
        wb = Workbook()
        sheets = CompositionFactory.from_model(self).sheets_from_model(self)

        first_sheet = True
        for sheet in sheets:
            if first_sheet:
                ws = wb.active
                first_sheet = False
            else:
                ws = wb.create_sheet()
            sheet.populate_worksheet(ws)
        return wb

    def to_file(self, path: Union[str, Path]):
        """Saves the model to a Excel file."""
        wb = self.workbook()
        wb.save(path)

    @classmethod
    def _property_keys(cls) -> List[str]:
        """Returns the names of the properties used in the excel header."""
        return [field.alias for field in cls.__fields__.values()]

    @classmethod
    def _ignore_key(cls, key: Optional[str]) -> bool:
        """
        Takes a property name (key) and checks if ignore_additional_rows is
        enabled the property is described by the model. This method is used by
        the from_workbook methods to determine if a cell should be added to the
        import dict.
        """
        ignore = getattr(cls.__config__, "ignore_additional_rows", None)
        if ignore is not None and not ignore:
            return False
        if key is None:
            return True
        return key not in cls._property_keys()

class XlsxConfig(BaseConfig):
    """
    Extends pydantic's config class with some Excel specific stuff.
    """

    header_font: Optional[Font] = Font(
        name="Arial",
        bold=True,
    )
    """
    Font of the header row (first row). Defaults to:
    `Font(name="Arial", bold=True)`.
    """
    font: Optional[Font] = None
    """Font for the non header rows. Defaults to `None`."""
    header_alignment: Optional[Alignment] = None
    """
    Optional alignment for the header cells. Defaults to `None`."""
    freeze_cell: Optional[str] = "A2"
    """
    Define a cell coordinate where the cells should be freeze. The same value
    is used to calculate the rows and columns which should be repeated on each
    printed page. Defaults to `A2` (header row (i.e. first row) stick to the
    top).

    Per default this value is also used to calculate
    `XlsxConfig.print_title_columns` and `XlsxConfig.print_title_rows` (rows
    and columns which should repeated on each printed page)
    """
    ignore_additional_columns: bool = False
    """
    When true, additional row exceeding the models definition will be ignored
    when importing data from a existing xlsx file. The presence of such rows
    will otherwise lead to an validation error of pydantic. Use this option
    with caution.
    """
    disable_width_calculation: bool = False
    """
    The library uses a primitive algorithm to set the width for each column
    based on number of chars of the longest cell content.
    """
    print_horizontal_centered: bool = True
    """
    Whether to horizontally center the content when printing the document.
    """
    print_vertical_centered: bool = True
    """
    Whether to vertically center the content when printing the document.
    """
    print_title_columns: Optional[str] = None
    """
    Defines the range of columns which should be repeated on each page in
    the print out. Per default this is calculated based on the value of
    `XlsxConfig.freeze_cell`. To disable the feature set this property to an
    empty string (`""`). You can define individual column ranges using the
    format `START_COLUMN:END_COLUMN`. E.g. `A:B` will repeat the first two
    columns.
    """
    print_title_rows: Optional[str] = None
    """
    Defines the range of rows which should be repeated on each page in the
    print out. Per default this is calculated based on the value of
    `XlsxConfig.freeze_cell`. To disable the feature set this property to
    an empty string (`""`). You can define individual row ranges using the
    format `START_ROW:END_ROW`. E.g. `1:2` will repeat the first two rows.
    """

    @classmethod
    def _print_title_columns(cls) -> Optional[str]:
        if cls.freeze_cell is None:
            return None
        freeze_column =\
            pyxl_utils.cell.coordinate_from_string(cls.freeze_cell)[0]
        max_column_num =\
            pyxl_utils.cell.column_index_from_string(freeze_column) - 1
        if max_column_num == 0:
            return None
        max_column = pyxl_utils.cell.get_column_letter(max_column_num)
        return f"A:{max_column}"

    @classmethod
    def _print_title_rows(cls) -> Optional[str]:
        if cls.print_title_rows is not None:
            return cls.print_title_rows
        if cls.freeze_cell is None:
            return None
        freeze_row =\
            pyxl_utils.cell.coordinate_from_string(cls.freeze_cell)[1] - 1
        if freeze_row == 0:
            return None
        return f"1:{freeze_row}"