From 6da6480edd74f87e5c2d6bd06ba1e08b3bf1ca42 Mon Sep 17 00:00:00 2001
From: c-bata <contact@c-bata.link>
Date: Thu, 13 Apr 2023 19:30:17 +0900
Subject: [PATCH] Introduce UserAttrFormWidget type

---
 optuna_dashboard/_form_widget.py     | 13 +++++--------
 optuna_dashboard/ts/types/index.d.ts | 20 ++++++++++++++------
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/optuna_dashboard/_form_widget.py b/optuna_dashboard/_form_widget.py
index e7e96d30..e78ee4fc 100644
--- a/optuna_dashboard/_form_widget.py
+++ b/optuna_dashboard/_form_widget.py
@@ -46,10 +46,7 @@ if TYPE_CHECKING:
         "TextInputWidgetJSON",
         {"type": Literal["text"], "description": Optional[str], "user_attr_key": Optional[str]},
     )
-    UserAttrRefJSON = TypedDict(
-        "UserAttrRefJSON",
-        {"type": Literal["user_attr"], "key": str, "user_attr_key": Optional[str]},
-    )
+    UserAttrRefJSON = TypedDict("UserAttrRefJSON", {"type": Literal["user_attr"], "key": str})
     FormWidgetJSON = TypedDict(
         "FormWidgetJSON",
         {
@@ -156,8 +153,6 @@ class ObjectiveUserAttrRef:
         return {
             "type": "user_attr",
             "key": self.key,
-            # Set 'user_attr_key' to simplify the frontend source code.
-            "user_attr_key": self.key,
         }
 
     @classmethod
@@ -202,7 +197,9 @@ def register_objective_form_widgets(
 ) -> None:
     if len(study.directions) != len(widgets):
         raise ValueError("The length of actions must be the same with the number of objectives.")
-    if any(not isinstance(w, ObjectiveUserAttrRef) and w.user_attr_key is not None for w in widgets):
+    if any(
+        not isinstance(w, ObjectiveUserAttrRef) and w.user_attr_key is not None for w in widgets
+    ):
         warnings.warn("`user_attr_key` specified, but it will not be used.")
     form_widgets: FormWidgetJSON = {
         "output_type": "objective",
@@ -218,7 +215,7 @@ def register_user_attr_form_widgets(
     widget_dicts: list[Union[ChoiceWidgetJSON, SliderWidgetJSON, TextInputWidgetJSON]] = []
     for w in widgets:
         if isinstance(w, ObjectiveUserAttrRef):
-            raise ValueError("ObjectiveUserAttrRef can't be specified in register_user_attr_form_widgets.")
+            raise ValueError("ObjectiveUserAttrRef can't be specified.")
         if w.user_attr_key is None:
             raise ValueError("`user_attr_key` is not specified.")
         user_attr_keys.add(w.user_attr_key)
diff --git a/optuna_dashboard/ts/types/index.d.ts b/optuna_dashboard/ts/types/index.d.ts
index 9e8b4456..24548a2c 100644
--- a/optuna_dashboard/ts/types/index.d.ts
+++ b/optuna_dashboard/ts/types/index.d.ts
@@ -157,20 +157,28 @@ type ObjectiveTextInputWidget = {
 type ObjectiveUserAttrRef = {
   type: "user_attr"
   key: string
-  user_attr_key?: string
 }
 
-// TODO(kenshin): Rename this type to FormWidget or something.
 type ObjectiveFormWidget =
   | ObjectiveChoiceWidget
   | ObjectiveSliderWidget
   | ObjectiveTextInputWidget
   | ObjectiveUserAttrRef
 
-type FormWidgets = {
-  output_type: "objective" | "user_attr"
-  widgets: ObjectiveFormWidget[]
-}
+type UserAttrFormWidget =
+  | ObjectiveChoiceWidget
+  | ObjectiveSliderWidget
+  | ObjectiveTextInputWidget
+
+type FormWidgets =
+  | {
+      output_type: "objective"
+      widgets: ObjectiveFormWidget[]
+    }
+  | {
+      output_type: "user_attr"
+      widgets: UserAttrFormWidget[]
+    }
 
 type StudyDetail = {
   id: number