From 917a62c9738b83d0f0cebc796d86633ac15d7a83 Mon Sep 17 00:00:00 2001
From: YooSunyoung <luysunyoung9@gmail.com>
Date: Tue, 10 Dec 2024 16:31:19 +0100
Subject: [PATCH] Optional widget fields setter/getter

---
 src/ess/reduce/widgets/_optional_widget.py | 23 +++++++++++++
 tests/widget_test.py                       | 40 +++++++++++++++++++++-
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/src/ess/reduce/widgets/_optional_widget.py b/src/ess/reduce/widgets/_optional_widget.py
index 26f5dd7e..6a33abf2 100644
--- a/src/ess/reduce/widgets/_optional_widget.py
+++ b/src/ess/reduce/widgets/_optional_widget.py
@@ -4,6 +4,7 @@
 
 from ipywidgets import HTML, HBox, Layout, RadioButtons, Widget
 
+from ._base import WidgetWithFieldsProtocol
 from ._config import default_style
 
 
@@ -64,3 +65,25 @@ def value(self, value: Any) -> None:
         else:
             self._option_box.value = self.name
             self.wrapped.value = value
+
+    def set_fields(self, new_values: Any) -> None:
+        # Set the value of the option box
+        if new_values is not None:
+            self._option_box.value = self.name
+        else:
+            self._option_box.value = None
+        # Set the value of the wrapped widget
+        if isinstance(self.wrapped, WidgetWithFieldsProtocol):
+            self.wrapped.set_fields(new_values)
+        elif new_values is not None:
+            self.wrapped.value = new_values
+        else:
+            ...  # Do not set the value of the wrapped widget
+
+    def get_fields(self) -> dict[str, Any] | None:
+        if self._option_box.value is None:
+            return None
+        elif isinstance(self.wrapped, WidgetWithFieldsProtocol):
+            return self.wrapped.get_fields()
+        else:
+            return self.wrapped.value
diff --git a/tests/widget_test.py b/tests/widget_test.py
index 793c4e9b..91e5959f 100644
--- a/tests/widget_test.py
+++ b/tests/widget_test.py
@@ -8,9 +8,15 @@
 import scipp as sc
 from ipywidgets import FloatText, IntText
 
-from ess.reduce.parameter import BinEdgesParameter, Parameter, parameter_registry
+from ess.reduce.parameter import (
+    BinEdgesParameter,
+    Parameter,
+    Vector3dParameter,
+    parameter_registry,
+)
 from ess.reduce.ui import WorkflowWidget, workflow_widget
 from ess.reduce.widgets import OptionalWidget, SwitchWidget, create_parameter_widget
+from ess.reduce.widgets._base import WidgetWithFieldsProtocol
 from ess.reduce.workflow import register_workflow, workflow_registry
 
 SwitchableInt = NewType('SwitchableInt', int)
@@ -166,6 +172,38 @@ def test_switchable_optional_parameter_switchable_first() -> None:
     assert isinstance(dummy_widget.wrapped, OptionalWidget)
 
 
+def test_optional_widget_set_value_get_fields() -> None:
+    optional_param = Parameter('a', 'a', 1, optional=True)
+    optional_widget = create_parameter_widget(optional_param)
+    assert isinstance(optional_widget, WidgetWithFieldsProtocol)
+    assert isinstance(optional_widget, OptionalWidget)
+    # Check initial state
+    assert optional_widget._option_box.value is None
+    assert optional_widget.get_fields() is None
+    # Update the value of the wrapped widget
+    optional_widget.value = 'test'
+    # Check the fields
+    assert optional_widget.get_fields() == 'test'
+
+
+def test_optional_widget_set_fields_get_fields() -> None:
+    optional_param = Vector3dParameter(
+        'a', 'a', sc.vector([1, 2, 3], unit='m'), optional=True
+    )
+    optional_widget = create_parameter_widget(optional_param)
+    assert isinstance(optional_widget, WidgetWithFieldsProtocol)
+    assert isinstance(optional_widget, OptionalWidget)
+    # Check initial state
+    assert optional_widget._option_box.value is None
+    assert optional_widget.get_fields() is None
+    # Update the value of the wrapped widget
+    optional_widget.set_fields({'x': 4, 'y': 5, 'z': 6, 'unit': 'm'})
+    assert optional_widget.value == sc.vector([4, 5, 6], unit='m')
+    # Check the fields and the option box value
+    assert optional_widget.get_fields() == {'x': 4, 'y': 5, 'z': 6, 'unit': 'm'}
+    assert optional_widget._option_box.value == optional_param.name
+
+
 def test_optional_widget_dispatch() -> None:
     optional_param = Parameter('a', 'a', 1, optional=True)
     assert isinstance(create_parameter_widget(optional_param), OptionalWidget)