examples/decoupling_io/adapters.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import dataclasses
from collections.abc import Collection
from os import PathLike
from typing import Any, Union

# This is not necessary once this PR gets merged: https://github.com/apache/hamilton/pull/467.
try:
    import sklearn.inspection
    import sklearn.metrics
    import sklearn.model_selection
except ImportError as e:
    raise NotImplementedError("scikit-learn is not installed.") from e


from hamilton import registry
from hamilton.io import utils
from hamilton.io.data_adapters import DataSaver

display_classes = [
    "ConfusionMatrixDisplay",
    "DetCurveDisplay",
    "PrecisionRecallDisplay",
    "PredictionErrorDisplay",
    "RocCurveDisplay",
    "CalibrationDisplay",
    "DecisionBoundaryDisplay",
    "LearningCurveDisplay",
    "PartialDependenceDisplay",
    "ValidationCurveDisplay",
]
SKLEARN_PLOT_TYPES = []
for class_name in display_classes:
    # get the attribute via string from sklearn.metrics; if not found return None
    class_metrics = getattr(sklearn.metrics, class_name, None)
    class_inspection = getattr(sklearn.inspection, class_name, None)
    class_model_selection = getattr(sklearn.model_selection, class_name, None)
    if class_metrics:
        SKLEARN_PLOT_TYPES.append(class_metrics)
    if class_inspection:
        SKLEARN_PLOT_TYPES.append(class_inspection)
    if class_model_selection:
        SKLEARN_PLOT_TYPES.append(class_model_selection)

SKLEARN_PLOT_TYPES_ANNOTATION = Union[tuple(SKLEARN_PLOT_TYPES)]


@dataclasses.dataclass
class SklearnPlotSaver(DataSaver):
    path: str | PathLike
    # kwargs
    dpi: float = 200
    format: str = "png"
    metadata: dict | None = None
    bbox_inches: str = None
    pad_inches: float = 0.1
    backend: str | None = None
    papertype: str = None
    transparent: bool = None
    bbox_extra_artists: list | None = None
    pil_kwargs: dict | None = None

    @classmethod
    def applicable_types(cls) -> Collection[type]:
        return SKLEARN_PLOT_TYPES

    def _get_saving_kwargs(self) -> dict[str, Any]:
        kwargs = {}
        if self.dpi is not None:
            kwargs["dpi"] = self.dpi
        if self.format is not None:
            kwargs["format"] = self.format
        if self.metadata is not None:
            kwargs["metadata"] = self.metadata
        if self.bbox_inches is not None:
            kwargs["bbox_inches"] = self.bbox_inches
        if self.pad_inches is not None:
            kwargs["pad_inches"] = self.pad_inches
        if self.backend is not None:
            kwargs["backend"] = self.backend
        if self.papertype is not None:
            kwargs["papertype"] = self.papertype
        if self.transparent is not None:
            kwargs["transparent"] = self.transparent
        if self.bbox_extra_artists is not None:
            kwargs["bbox_extra_artists"] = self.bbox_extra_artists
        if self.pil_kwargs is not None:
            kwargs["pil_kwargs"] = self.pil_kwargs
        return kwargs

    def save_data(self, data: SKLEARN_PLOT_TYPES_ANNOTATION) -> dict[str, Any]:
        data.plot()
        data.figure_.savefig(self.path, **self._get_saving_kwargs())
        return utils.get_file_metadata(self.path)

    @classmethod
    def name(cls) -> str:
        return "png"


def register_data_loaders():
    """Function to register the data loaders for this extension."""
    registry.register_adapter(SklearnPlotSaver)


register_data_loaders()