tasks.py

"""Functionality around tasks

Tasks are used to implement "undo" and "redo" functionality.

"""
from __future__ import annotations

import shutil
from pathlib import Path
from tempfile import mkdtemp
from uuid import uuid4

from skops import card
from skops.card._model_card import PlotSection, split_subsection_names
from streamlit.runtime.uploaded_file_manager import UploadedFile


class Task:
    """(Abstract) base class for tasks"""

    def do(self) -> None:
        raise NotImplementedError

    def undo(self) -> None:
        raise NotImplementedError


class TaskState:
    """Tracking the state of tasks"""

    def __init__(self) -> None:
        self.done_list: list[Task] = []
        self.undone_list: list[Task] = []

    def undo(self) -> None:
        if not self.done_list:
            return

        task = self.done_list.pop(-1)
        task.undo()
        self.undone_list.append(task)

    def redo(self) -> None:
        if not self.undone_list:
            return

        task = self.undone_list.pop(-1)
        task.do()
        self.done_list.append(task)

    def add(self, task: Task) -> None:
        task.do()
        self.done_list.append(task)
        self.undone_list.clear()

    def reset(self) -> None:
        self.done_list.clear()
        self.undone_list.clear()


class AddSectionTask(Task):
    """Add a new text section"""

    def __init__(
        self,
        model_card: card.Card,
        title: str,
        content: str,
    ) -> None:
        self.model_card = model_card
        self.title = title
        self.key = title + " " + str(uuid4())[:6]
        self.content = content

    def do(self) -> None:
        self.model_card.add(**{self.key: self.content})
        section = self.model_card.select(self.key)
        section.title = split_subsection_names(self.title)[-1]

    def undo(self) -> None:
        self.model_card.delete(self.key)


class AddFigureTask(Task):
    """Add a new figure section

    Figure always starts out with dummy image cat.png.

    """

    def __init__(
        self,
        model_card: card.Card,
        path: Path,
        title: str,
        content: str,
    ) -> None:
        self.model_card = model_card
        self.title = title

        # Create a unique file name, since the same image can exist more than
        # once per model card.
        fname = Path(content)
        stem = fname.stem
        suffix = fname.suffix
        uniq = str(uuid4())[:6]
        new_fname = str(path / stem) + "_" + uniq + suffix

        self.key = title + " " + uniq
        self.content = Path(new_fname)

    def do(self) -> None:
        shutil.copy("cat.png", self.content)
        self.model_card.add_plot(**{self.key: self.content})
        section = self.model_card.select(self.key)
        section.title = split_subsection_names(self.title)[-1]
        section.is_fig = True  # type: ignore

    def undo(self) -> None:
        self.content.unlink(missing_ok=True)
        self.model_card.delete(self.key)


class DeleteSectionTask(Task):
    """Delete a section

    The section is not completely removed from the underlying data structure,
    but only turned invisible.

    """

    def __init__(
        self,
        model_card: card.Card,
        key: str,
        path: Path | None,
    ) -> None:
        self.model_card = model_card
        self.key = key
        # when 'deleting' a file, move it to a temp file
        self.path = path
        self.tmp_path = Path(mkdtemp(prefix="skops-")) / str(uuid4())

    def do(self) -> None:
        self.model_card.select(self.key).visible = False
        if self.path:
            shutil.move(self.path, self.tmp_path)

    def undo(self) -> None:
        self.model_card.select(self.key).visible = True
        if self.path:
            shutil.move(self.tmp_path, self.path)


class UpdateSectionTask(Task):
    """Change the title or content of a text section"""

    def __init__(
        self,
        model_card: card.Card,
        key: str,
        old_name: str,
        new_name: str,
        old_content: str,
        new_content: str,
    ) -> None:
        self.model_card = model_card
        self.key = key
        self.old_name = old_name
        self.new_name = new_name
        self.old_content = old_content
        self.new_content = new_content

    def do(self) -> None:
        section = self.model_card.select(self.key)
        new_title = split_subsection_names(self.new_name)[-1]
        section.title = new_title
        section.content = self.new_content

    def undo(self) -> None:
        section = self.model_card.select(self.key)
        old_title = split_subsection_names(self.old_name)[-1]
        section.title = old_title
        section.content = self.old_content


class UpdateFigureTask(Task):
    """Change the title or image of a figure section

    Changing the title is easy, just replace it and be done.

    Changing the figure is a bit more tricky. The old figure is in the hf_path
    under its old name. The new figure is an UploadFile object. For the DO
    operation, move the old figure to a temporary file and store the UploadFile
    content to a new file (which may have a different name).

    For the UNDO operation, delete the new figure (its content is still stored
    in the UploadFile) and move back the old figure from its temporary file to
    the original location (with its original name).

    """

    def __init__(
        self,
        model_card: card.Card,
        key: str,
        old_name: str,
        new_name: str,
        data: UploadedFile | None,
        new_path: Path | None,
        old_path: Path | None,
    ) -> None:
        self.model_card = model_card
        self.key = key
        self.old_name = old_name
        self.new_name = new_name
        self.old_data = self.model_card.select(self.key).content
        self.new_path = new_path
        self.old_path = old_path
        # when 'deleting' the old image, move to temp path
        self.tmp_path = Path(mkdtemp(prefix="skops-")) / str(uuid4())

        if not data:
            self.new_data = self.old_data
        else:
            self.new_data = data

    def do(self) -> None:
        section = self.model_card.select(self.key)
        new_title = split_subsection_names(self.new_name)[-1]
        section.title = self.title = new_title
        if self.new_data == self.old_data:  # image is same
            return

        # write figure
        # note: this can still be the same image if the image is a file, there
        # is no test to check, e.g., the hash of the image
        shutil.move(self.old_path, self.tmp_path)

        with open(self.new_path, "wb") as f:
            f.write(self.new_data.getvalue())
        section.content = PlotSection(
            alt_text=self.new_data.name,
            path=self.new_path,
        )

    def undo(self) -> None:
        section = self.model_card.select(self.key)
        old_title = split_subsection_names(self.old_name)[-1]
        section.title = old_title
        if self.new_data == self.old_data:  # image is same
            return

        self.new_path.unlink(missing_ok=True)
        shutil.move(self.tmp_path, self.old_path)
        section.content = self.old_data


class AddMetricsTask(Task):
    """Add new metrics"""

    def __init__(
        self,
        model_card: card.Card,
        metrics: dict[str, str | int | float],
    ) -> None:
        self.model_card = model_card
        self.old_metrics = model_card._metrics.copy()
        self.new_metrics = metrics

    def do(self) -> None:
        self.model_card._metrics.clear()
        self.model_card.add_metrics(**self.new_metrics)

    def undo(self) -> None:
        self.model_card._metrics.clear()
        self.model_card.add_metrics(**self.old_metrics)