Model/epc_data/attributes/FloorAttributes.py

from typing import Dict, Union, Optional
from epc_data.attributes.attribute_utils import extract_thermal_transmittence, search_split_description


class FloorAttributes:

    def __init__(self, description: str):
        """
        Constructor for CleanFloor class.

        :param description: Description to be cleaned.
        """
        self.description: str = description

    def clean(self) -> Dict[str, Union[str, bool, int, None]]:
        """
        Cleans the given description based on predefined rules.

        :return: Cleaned data in a dictionary format.
        """

        description_lower = self.description.lower().strip()

        assumed: bool = "assumed" in description_lower
        to_unheated_space: bool = "to unheated space" in description_lower
        to_external_air: bool = "to external air" in description_lower
        is_suspended: bool = "suspended" in description_lower
        is_solid: bool = "solid" in description_lower
        thermal_transmittence: Optional[float] = None
        thermal_transmittence_unit: Optional[str] = None
        insulation_thickness: Optional[str] = None
        is_valid: bool = "invalid" not in description_lower

        if "another dwelling below" in description_lower or "other premises below" in description_lower:
            return self._make_clean_output(
                is_valid=is_valid,
                has_dwelling_below=True,
                thermal_transmittence=None,
                thermal_transmittence_unit=None,
                to_unheated_space=to_unheated_space,
                to_external_air=to_external_air,
                is_suspended=is_suspended,
                is_solid=is_solid,
                assumed=assumed,
                insulation_thickness=insulation_thickness
            )

        if "insulation" in description_lower or "insulated" in description_lower:
            insulation_thickness = self._find_insulation_thickness(
                description_lower, to_unheated_space, to_external_air, is_suspended, is_solid
            )
        elif "thermal transmittance" in description_lower:
            thermal_transmittence, thermal_transmittence_unit = extract_thermal_transmittence(description_lower)

        return self._make_clean_output(
            is_valid=is_valid,
            has_dwelling_below=False,
            thermal_transmittence=thermal_transmittence,
            thermal_transmittence_unit=thermal_transmittence_unit,
            to_unheated_space=to_unheated_space,
            to_external_air=to_external_air,
            is_suspended=is_suspended,
            is_solid=is_solid,
            assumed=assumed,
            insulation_thickness=insulation_thickness
        )

    @classmethod
    def _find_insulation_thickness(
        cls,
        description_lower: str,
        to_unheated_space: bool,
        to_external_air: bool,
        is_suspended: bool,
        is_solid: bool
    ) -> str:
        """
        Method to find the insulation thickness from the description.

        :param description_lower: Description of the floor in lower case.
        :param to_unheated_space: Boolean indicating if there's an unheated space below.
        :param to_external_air: Boolean indicating if the floor is exposed to external air.
        :param is_suspended: Boolean indicating if the floor is suspended.
        :param is_solid: Boolean indicating if the floor is solid.
        :return: Insulation thickness as a string.
        """
        if to_unheated_space | to_external_air | is_suspended | is_solid:
            if to_unheated_space:
                split_on = "to unheated space,"
            elif to_external_air:
                split_on = "to external air,"
            elif is_suspended:
                split_on = "suspended,"
            elif is_solid:
                split_on = "solid,"
            else:
                raise NotImplementedError("Not handled this 2")
            desc_split = description_lower.split(split_on)[-1].strip().split("(assumed)")[0].strip()
            return search_split_description(desc_split)

        raise NotImplementedError("Not handled this")

    @staticmethod
    def _make_clean_output(
        is_valid: bool,
        has_dwelling_below: bool,
        thermal_transmittence: Optional[float],
        thermal_transmittence_unit: Optional[str],
        to_unheated_space: bool,
        to_external_air: bool,
        is_suspended: bool,
        is_solid: bool,
        assumed: bool,
        insulation_thickness: Optional[str]
    ) -> Dict[str, Union[str, bool, int, None]]:
        """
        Creates a dictionary with the cleaned data.

        :param is_valid: Boolean indicating if the description is valid.
        :param has_dwelling_below: Boolean indicating if there's a dwelling below.
        :param thermal_transmittence: Thermal transmittence value.
        :param thermal_transmittence_unit: Unit of thermal transmittence.
        :param to_unheated_space: Boolean indicating if there's an unheated space below.
        :param to_external_air: Boolean indicating if the floor is exposed to external air.
        :param is_suspended: Boolean indicating if the floor is suspended.
        :param is_solid: Boolean indicating if the floor is solid.
        :param assumed: Boolean indicating if the data was assumed.
        :param insulation_thickness: Insulation thickness as a string.
        :return: Dictionary with cleaned data.
        """
        return {
            "is_valid": is_valid,
            "has_dwelling_below": has_dwelling_below,
            "thermal_transmittence": thermal_transmittence,
            "thermal_transmittence_unit": thermal_transmittence_unit,
            "to_unheated_space": to_unheated_space,
            "to_external_air": to_external_air,
            "is_suspended": is_suspended,
            "is_solid": is_solid,
            "assumed": assumed,
            "insulation_thickness": insulation_thickness
        }