From 1366936472cc0d30bc10853a13e72058cb4ee300 Mon Sep 17 00:00:00 2001
From: c-bata <contact@c-bata.link>
Date: Wed, 4 Jan 2023 16:28:30 +0900
Subject: [PATCH] Add Python API to save the note

---
 README.md                    |  8 +++++
 optuna_dashboard/__init__.py |  2 ++
 optuna_dashboard/_note.py    | 62 ++++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+)

diff --git a/README.md b/README.md
index 26310af4..0853df42 100644
--- a/README.md
+++ b/README.md
@@ -88,6 +88,14 @@ This function uses wsgiref module which is not intended for the production use.
 This function exposes WSGI interface for people who want to run on the
 production-class WSGI servers like Gunicorn or uWSGI.
 
+**`save_study_note(study: Study, body: string) -> None`**
+
+Save the note (Markdown format) to the Study.
+
+**`save_trial_note(trial: Trial, body: string) -> None`**
+
+Save the note (Markdown format) to the Trial.
+
 </details>
 
 ## Using an official Docker image
diff --git a/optuna_dashboard/__init__.py b/optuna_dashboard/__init__.py
index 1d3b4197..9c23bd5e 100644
--- a/optuna_dashboard/__init__.py
+++ b/optuna_dashboard/__init__.py
@@ -1,5 +1,7 @@
 from ._app import run_server  # noqa
 from ._app import wsgi  # noqa
+from ._note import save_study_note  # noqa
+from ._note import save_trial_note  # noqa
 
 
 __version__ = "0.9.0b1"
diff --git a/optuna_dashboard/_note.py b/optuna_dashboard/_note.py
index 4e872e1e..1ed807dd 100644
--- a/optuna_dashboard/_note.py
+++ b/optuna_dashboard/_note.py
@@ -4,6 +4,7 @@ import math
 from typing import Any
 from typing import TYPE_CHECKING
 
+import optuna
 from optuna.storages import BaseStorage
 
 
@@ -22,6 +23,67 @@ if TYPE_CHECKING:
 SYSTEM_ATTR_MAX_LENGTH = 2045
 
 
+def save_study_note(study: optuna.Study, body: str) -> None:
+    """Save the note (Markdown format) to the Study.
+
+    Example:
+
+       .. code-block:: python
+
+          import optuna
+          from optuna_dashboard import save_study_note
+
+          study = optuna.create_study()
+
+          note = textwrap.dedent('''\
+          ## Hello
+
+          You can *freely* take a **note** that is associated with the study.
+          ''')
+          save_study_note(study, note)
+
+    """
+    storage = study._storage
+    study_id = study._study_id
+    system_attrs = storage.get_study_system_attrs(study_id)
+    next_ver = system_attrs.get(note_ver_key(None), 0) + 1
+    save_note(storage, study_id, None, next_ver, body)
+
+
+def save_trial_note(trial: optuna.Trial, body: str) -> None:
+    """Save the note (Markdown format) to the Trial.
+
+    Example:
+
+       .. code-block:: python
+
+          import optuna
+          import textwrap
+          from optuna_dashboard import save_trial_note
+
+          def objective_single(trial: optuna.Trial) -> float:
+              x1 = trial.suggest_float("x1", 0, 10)
+              x2 = trial.suggest_float("x2", 0, 10)
+
+              note = textwrap.dedent(f'''\
+              ## Trial {trial._trial_id}
+
+              $$
+              y = (x1 - 2)^{{2}} + (x2 - 5)^{{2}} = ({x1} - 2)^{{2}} + ({x2} - 5)^{{2}}
+              $$
+              ''')
+              save_trial_note(trial, note)
+              return (x1 - 2) ** 2 + (x2 - 5) ** 2
+    """
+    storage = trial.storage
+    trial_id = trial._trial_id
+    study_id = trial.study._study_id
+
+    system_attrs = storage.get_study_system_attrs(study_id)
+    next_ver = system_attrs.get(note_ver_key(trial_id), 0) + 1
+    save_note(storage, study_id, trial_id, next_ver, body)
+
+
 def note_ver_key(trial_id: Optional[int]) -> str:
     prefix = "dashboard:note_ver"
     if trial_id is None: