Source code for openstef_models.integrations.joblib.joblib_model_serializer

# SPDX-FileCopyrightText: 2025 Contributors to the OpenSTEF project <openstef@lfenergy.org>
#
# SPDX-License-Identifier: MPL-2.0
"""Local model storage implementation using joblib serialization.

Provides file-based persistence for ForecastingModel instances using joblib's
pickle-based serialization. This storage backend is suitable for development,
testing, and single-machine deployments where models need to be persisted
to the local filesystem.
"""

from typing import BinaryIO, ClassVar, override

from openstef_core.exceptions import MissingExtraError
from openstef_models.mixins.model_serializer import ModelSerializer

try:
    import joblib
except ImportError as e:
    raise MissingExtraError("joblib", package="openstef-models") from e




class JoblibModelSerializer(ModelSerializer):
    """File-based model storage using joblib serialization.

    Provides persistent storage for ForecastingModel instances on the local
    filesystem. Models are serialized using joblib and stored as pickle files
    in the specified directory.

    This storage implementation is suitable for development, testing, and
    single-machine deployments where simple file-based persistence is sufficient.

    Note:
        joblib.dump() and joblib.load() are based on the Python pickle serialization model,
        which means that arbitrary Python code can be executed when loading a serialized object
        with joblib.load().

        joblib.load() should therefore never be used to load objects from an untrusted source
        or otherwise you will introduce a security vulnerability in your program.

    Invariants:
        - Models are stored as .pkl files in the configured storage directory
        - Model files use the pattern: {model_id}.pkl
        - Storage directory is created automatically if it doesn't exist
        - Load operations fail with ModelNotFoundError if model file doesn't exist

    Example:
        Basic usage with model persistence

        >>> from pathlib import Path
        >>> from openstef_models.models.forecasting_model import ForecastingModel
        >>> storage = LocalModelStorage(storage_dir=Path("./models"))  # doctest: +SKIP
        >>> storage.save_model("my_model", my_forecasting_model)  # doctest: +SKIP
        >>> loaded_model = storage.load_model("my_model")  # doctest: +SKIP
    """

    extension: ClassVar[str] = "joblib"



    @override
    def serialize(self, model: object, file: BinaryIO) -> None:
        joblib.dump(model, file)  # type: ignore[reportUnknownMemberType]




    @override
    def deserialize(self, file: BinaryIO) -> object:
        return joblib.load(file)  # type: ignore[reportUnknownMemberType]




__all__ = ["JoblibModelSerializer"]