from __future__ import annotations from pathlib import Path 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: def __init__(self, model_card: card.Card) -> None: self.model_card = model_card def do(self) -> None: raise NotImplementedError def undo(self) -> None: raise NotImplementedError class TaskState: 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): 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): 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_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.model_card.delete(self.key) class DeleteSectionTask(Task): def __init__( self, model_card: card.Card, key: str, ) -> None: self.model_card = model_card self.key = key def do(self) -> None: self.model_card.select(self.key).visible = False def undo(self) -> None: self.model_card.select(self.key).visible = True class UpdateSectionTask(Task): 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): def __init__( self, model_card: card.Card, key: str, old_name: str, new_name: str, data: UploadedFile | None, 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.path = path 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 with open(self.path, "wb") as f: f.write(self.new_data.getvalue()) section.content = PlotSection( alt_text=self.new_data.name, path=self.path, ).format() 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.path.unlink(missing_ok=True) section.content = self.old_data