Dataloadingplan

@abstractmethod
def apply(self, *args, **kwargs):
    """Abstract method representing an application of the DataLoadingBlock"""
    pass

def deserialize(self, load_from: dict) -> TDataLoadingBlock:
    """Reconstruct the DataLoadingBlock from a serialized version.

    Args:
        load_from (dict): a dictionary as obtained by the serialize function.
    Returns:
        the self instance
    """
    self._serialization_validator.validate(
        load_from, FedbiomedLoadingBlockValueError
    )
    self.__serialization_id = load_from["dlb_id"]
    return self

def get_serialization_id(self):
    """Expose serialization id as read-only"""
    return self.__serialization_id

@staticmethod
def instantiate_class(loading_block: dict) -> TDataLoadingBlock:
    """Instantiate one [DataLoadingBlock][fedbiomed.common.dataloadingplan.DataLoadingBlock]
    object of the type defined in the arguments.

    Uses the `loading_block_module` and `loading_block_class` fields of the loading_block argument to
    identify the type of [DataLoadingBlock][fedbiomed.common.dataloadingplan.DataLoadingBlock]
    to be instantiated, then calls its default constructor.
    Note that this function **does not call deserialize**.

    Args:
        loading_block (dict): [DataLoadingBlock][fedbiomed.common.dataloadingplan.DataLoadingBlock]
            metadata in the format returned by the serialize function.
    Returns:
        A default-constructed instance of a
            [DataLoadingBlock][fedbiomed.common.dataloadingplan.DataLoadingBlock]
            of the type defined in the metadata.
    Raises:
       FedbiomedLoadingBlockError: if the instantiation process raised any exception.
    """
    try:
        dlb_module = import_module(loading_block["loading_block_module"])  # noqa: F841
        dlb = eval(f"dlb_module.{loading_block['loading_block_class']}()")
    except Exception as e:
        msg = (
            f"{ErrorNumbers.FB614.value}: could not instantiate DataLoadingBlock from the following metadata: "
            + f"{loading_block} because of {type(e).__name__}: {e}"
        )
        logger.debug(msg)
        raise FedbiomedLoadingBlockError(msg) from e
    return dlb

@staticmethod
def instantiate_key(
    key_module: str, key_classname: str, loading_block_key_str: str
) -> DataLoadingBlockTypes:
    """Imports and loads [DataLoadingBlockTypes][fedbiomed.common.constants.DataLoadingBlockTypes]
    regarding the passed arguments

    Args:
        key_module (str): _description_
        key_classname (str): _description_
        loading_block_key_str (str): _description_

    Raises:
        FedbiomedDataLoadingPlanError: _description_

    Returns:
        DataLoadingBlockTypes: _description_
    """
    try:
        keys = import_module(key_module)  # noqa: F841
        loading_block_key = eval(f"keys.{key_classname}('{loading_block_key_str}')")
    except Exception as e:
        msg = (
            f"{ErrorNumbers.FB615.value} Error deserializing loading block key "
            + f"{loading_block_key_str} with path {key_module}.{key_classname} "
            + f"because of {type(e).__name__}: {e}"
        )
        logger.debug(msg)
        raise FedbiomedDataLoadingPlanError(msg) from e
    return loading_block_key

def serialize(self) -> dict:
    """Serializes the class in a format similar to json.

    Returns:
        a dictionary of key-value pairs sufficient for reconstructing
        the DataLoadingBlock.
    """
    return dict(
        loading_block_class=self.__class__.__qualname__,
        loading_block_module=self.__module__,
        dlb_id=self.__serialization_id,
    )

def deserialize(
    self, serialized_dlp: dict, serialized_loading_blocks: List[dict]
) -> TDataLoadingPlan:
    """Reconstruct the DataLoadingPlan][fedbiomed.common.dataloadingplan.DataLoadingPlan] from a serialized version.

    !!! warning "Calling this function will *clear* the contained [DataLoadingBlockTypes]."
        This function may not be used to "update" nor to "append to"
        a [DataLoadingPlan][fedbiomed.common.dataloadingplan.DataLoadingPlan].

    Args:
        serialized_dlp: a dictionary of data loading plan metadata, as obtained from the first output of the
            serialize function
        serialized_loading_blocks: a list of dictionaries of loading_block metadata, as obtained from the
            second output of the serialize function
    Returns:
        the self instance
    """
    self._serialization_validation.validate(
        serialized_dlp, FedbiomedDataLoadingPlanValueError
    )

    self.clear()
    self.dlp_id = serialized_dlp["dlp_id"]
    self.desc = serialized_dlp["dlp_name"]
    self.target_dataset_type = DatasetTypes(serialized_dlp["target_dataset_type"])
    for loading_block_key_str, dlb_id in serialized_dlp["loading_blocks"].items():
        key_module, key_classname = serialized_dlp["key_paths"][
            loading_block_key_str
        ]
        loading_block_key = DataLoadingBlock.instantiate_key(
            key_module, key_classname, loading_block_key_str
        )
        loading_block = next(
            filter(lambda x: x["dlb_id"] == dlb_id, serialized_loading_blocks)
        )
        dlb = DataLoadingBlock.instantiate_class(loading_block)
        self[loading_block_key] = dlb.deserialize(loading_block)
    return self

@staticmethod
def infer_dataset_type(dataset: Any) -> DatasetTypes:
    """Infer the type of a given dataset.

    This function provides the mapping between a dataset's class and the DatasetTypes enum. If the dataset exposes
    the correct interface (i.e. the get_dataset_type method) then it directly calls that, otherwise it tries to
    apply some heuristics to guess the type of dataset.

    Args:
        dataset: the dataset whose type we want to infer.
    Returns:
        a DatasetTypes enum element which identifies the type of the dataset.
    Raises:
        FedbiomedDataLoadingPlanValueError: if the dataset does not have a `get_dataset_type` method and moreover
            the type could not be guessed.
    """
    if hasattr(dataset, "get_dataset_type"):
        return dataset.get_dataset_type()
    elif dataset.__class__.__name__ == "ImageFolder":
        # ImageFolder could be both an images type or mednist. Try to identify mednist with some heuristic.
        if hasattr(dataset, "classes") and all(
            [
                x in dataset.classes
                for x in [
                    "AbdomenCT",
                    "BreastMRI",
                    "CXR",
                    "ChestCT",
                    "Hand",
                    "HeadCT",
                ]
            ]
        ):
            return DatasetTypes.MEDNIST
        else:
            return DatasetTypes.IMAGES
    elif dataset.__class__.__name__ == "MNIST":
        return DatasetTypes.DEFAULT
    msg = (
        f"{ErrorNumbers.FB615.value} Trying to infer dataset type of {dataset} is not supported "
        + f"for datasets of type {dataset.__class__.__qualname__}"
    )
    logger.debug(msg)
    raise FedbiomedDataLoadingPlanValueError(msg)

def serialize(self) -> Tuple[dict, List]:
    """Serializes the class in a format similar to json.

    Returns:
        a tuple sufficient for reconstructing the DataLoading plan. It includes:
            - a dictionary of key-value pairs with the
            [DataLoadingPlan][fedbiomed.common.dataloadingplan.DataLoadingPlan] parameters.
            - a list of dict containing the data for reconstruction all the DataLoadingBlock
                of the [DataLoadingPlan][fedbiomed.common.dataloadingplan.DataLoadingPlan]
    """
    return dict(
        dlp_id=self.dlp_id,
        dlp_name=self.desc,
        target_dataset_type=self.target_dataset_type.value,
        loading_blocks={
            key.value: dlb.get_serialization_id() for key, dlb in self.items()
        },
        key_paths={
            key.value: (f"{key.__module__}", f"{key.__class__.__qualname__}")
            for key in self.keys()
        },
    ), [dlb.serialize() for dlb in self.values()]

def apply_dlb(
    self,
    default_ret_value: Any,
    dlb_key: DataLoadingBlockTypes,
    *args: Optional[Any],
    **kwargs: Optional[Any],
) -> Any:
    """Apply one DataLoadingBlock identified by its key.

    Note that we want to easily support the case where the DataLoadingPlan
    is not activated, or the requested loading block is not contained in the
    DataLoadingPlan. This is achieved by providing a default return value
    to be returned when the above conditions are met. Hence, most of the
    calls to apply_dlb will look like this:
    ```
    value = self.apply_dlb(value, 'my-loading-block', my_apply_args)
    ```
    This will ensure that value is not changed if the DataLoadingPlan is
    not active.

    Args:
        default_ret_value: the value to be returned in case that the dlp
            functionality is not required
        dlb_key: the key of the DataLoadingBlock to be applied
        *args: forwarded to the DataLoadingBlock's apply function
        **kwargs: forwarded to the DataLoadingBlock's apply function
    Returns:
        the output of the DataLoadingBlock's apply function, or
            the default_ret_value when dlp is None or it does not contain
            the requested loading block
    """
    if not isinstance(dlb_key, DataLoadingBlockTypes):
        raise FedbiomedDataLoadingPlanValueError(
            f"Key {dlb_key} is not of enum type DataLoadingBlockTypes"
            f" in DataLoadingPlanMixin.apply_dlb"
        )
    if self._dlp is not None and dlb_key in self._dlp:
        return self._dlp[dlb_key].apply(*args, **kwargs)
    else:
        return default_ret_value

def clear_dlp(self):
    self._dlp = None

def set_dlp(self, dlp: DataLoadingPlan):
    """Sets the dlp if the target dataset type is appropriate"""
    if not isinstance(dlp, DataLoadingPlan):
        msg = (
            f"{ErrorNumbers.FB615.value} Trying to set a DataLoadingPlan but the argument is of type "
            + f"{type(dlp).__name__}"
        )
        logger.debug(msg)
        raise FedbiomedDataLoadingPlanValueError(msg)

    dataset_type = DataLoadingPlan.infer_dataset_type(
        self
    )  # `self` here will refer to the Dataset instance
    if (
        dlp.target_dataset_type != DatasetTypes.NONE
        and dataset_type != dlp.target_dataset_type
    ):
        raise FedbiomedDataLoadingPlanValueError(
            f"Trying to set {dlp} on dataset of type {dataset_type.value} but "
            f"the target type is {dlp.target_dataset_type}"
        )
    elif dlp.target_dataset_type == DatasetTypes.NONE:
        dlp.target_dataset_type = dataset_type
    self._dlp = dlp

def apply(self, key):
    """Returns the value mapped to the key, if it exists.

    Raises:
        FedbiomedLoadingBlockError: if map is not a dict or the key does not exist.
    """
    if not isinstance(self.map, dict) or key not in self.map:
        msg = f"{ErrorNumbers.FB614.value} Mapper block error: no key '{key}' in mapping dictionary"
        logger.debug(msg)
        raise FedbiomedLoadingBlockError(msg)
    return self.map[key]

def deserialize(self, load_from: dict) -> DataLoadingBlock:
    """Reconstruct the [DataLoadingBlock][fedbiomed.common.dataloadingplan.DataLoadingBlock]
    from a serialized version.

    Args:
        load_from (dict): a dictionary as obtained by the serialize function.
    Returns:
        the self instance
    """
    super(MapperBlock, self).deserialize(load_from)
    self.map = load_from["map"]
    return self

def serialize(self) -> dict:
    """Serializes the class in a format similar to json.

    Returns:
        a dictionary of key-value pairs sufficient for reconstructing
        the [DataLoadingBlock][fedbiomed.common.dataloadingplan.DataLoadingBlock].
    """
    ret = super(MapperBlock, self).serialize()
    ret.update({"map": self.map})
    return ret

@classmethod
def dlb_default_scheme(cls) -> Dict:
    """The dictionary of default validation rules for a serialized [DataLoadingBlock]."""
    return {
        "loading_block_class": {
            "rules": [str, cls._identifier_validation_hook],
            "required": True,
        },
        "loading_block_module": {
            "rules": [str, cls._identifier_validation_hook],
            "required": True,
        },
        "dlb_id": {
            "rules": [str, cls._serial_id_validation_hook],
            "required": True,
        },
    }

@classmethod
def dlp_default_scheme(cls) -> Dict:
    """The dictionary of default validation rules for a serialized [DataLoadingPlan]."""
    return {
        "dlp_id": {
            "rules": [str],
            "required": True,
        },
        "dlp_name": {
            "rules": [str],
            "required": True,
        },
        "target_dataset_type": {
            "rules": [str, cls._target_dataset_type_validator],
            "required": True,
        },
        "loading_blocks": {
            "rules": [dict, cls._loading_blocks_types_validator],
            "required": True,
        },
        "key_paths": {"rules": [dict, cls._key_paths_validator], "required": True},
    }

def update_validation_scheme(self, new_scheme: dict) -> None:
    """Updates the validation scheme.

    Args:
        new_scheme: (dict) new dict of rules
    """
    self._validation_scheme.update(new_scheme)

def validate(
    self,
    dlb_metadata: Dict,
    exception_type: Type[FedbiomedError],
    only_required: bool = True,
) -> None:
    """Validate a dict of dlb_metadata according to the _validation_scheme.

    Args:
        dlb_metadata (dict) : the [DataLoadingBlock] metadata, as returned by serialize or as loaded from the
            node database.
        exception_type (Type[FedbiomedError]): the type of the exception to be raised when validation fails.
        only_required (bool) : see SchemeValidator.populate_with_defaults
    Raises:
        exception_type: if the validation fails.
    """
    try:
        sc = SchemeValidator(self._validation_scheme)
    except RuleError as e:
        msg = ErrorNumbers.FB614.value + f": {e}"
        logger.critical(msg)
        raise exception_type(msg) from e

    try:
        dlb_metadata = sc.populate_with_defaults(
            dlb_metadata, only_required=only_required
        )
    except ValidatorError as e:
        msg = ErrorNumbers.FB614.value + f": {e}"
        logger.critical(msg)
        raise exception_type(msg) from e

    try:
        sc.validate(dlb_metadata)
    except ValidateError as e:
        msg = ErrorNumbers.FB614.value + f": {e}"
        logger.critical(msg)
        raise exception_type(msg) from e