From a8ff0b693a22d29956db6287707f6131fda69583 Mon Sep 17 00:00:00 2001
From: Joe Numainville <josephnumainville@deephaven.io>
Date: Thu, 18 Apr 2024 17:15:16 -0500
Subject: [PATCH] refactoring so it is still usable

---
 .../ui/src/deephaven/ui/_internal/utils.py    | 28 ++++++++++++-
 .../src/deephaven/ui/components/list_view.py  | 12 +++++-
 .../ui/src/deephaven/ui/components/picker.py  | 12 +++++-
 plugins/ui/test/deephaven/ui/test_utils.py    | 41 +++++++++++++++++++
 4 files changed, 90 insertions(+), 3 deletions(-)

diff --git a/plugins/ui/src/deephaven/ui/_internal/utils.py b/plugins/ui/src/deephaven/ui/_internal/utils.py
index 1cd6b66e3..27ea3d7d2 100644
--- a/plugins/ui/src/deephaven/ui/_internal/utils.py
+++ b/plugins/ui/src/deephaven/ui/_internal/utils.py
@@ -184,7 +184,7 @@ def wrap_callable(func: Callable) -> Callable:
         return func
 
 
-def create_props(args: dict[str, Any]) -> tuple[tuple[Any], dict[str, Any]]:
+def create_props(args: dict[str, Any]) -> tuple[tuple[Any, ...], dict[str, Any]]:
     """
     Create props from the args. Combines the named props with the kwargs props.
 
@@ -197,3 +197,29 @@ def create_props(args: dict[str, Any]) -> tuple[tuple[Any], dict[str, Any]]:
     children, props = args.pop("children"), args.pop("props")
     props.update(args)
     return children, props
+
+
+def unpack_item_table_source(
+    children: tuple[Any, ...],
+    props: dict[str, Any],
+    supported_args: set[str],
+) -> tuple[tuple[Any, ...], dict[str, Any]]:
+    """
+    Unpack children and props if the children are of type dict
+    and merge the supported arguments into the props.
+
+    Args:
+        children: The children to possibly unpack.
+        props: The props to unpack.
+        supported_args: The supported arguments for the ItemTableSource.
+
+    Returns:
+        The unpacked children and props.
+    """
+    if len(children) == 1 and isinstance(children[0], dict):
+        item_table_source = children[0]
+        children = (item_table_source.pop("table"),)
+        for key in supported_args:
+            if key in item_table_source:
+                props[key] = item_table_source.pop(key)
+    return children, props
diff --git a/plugins/ui/src/deephaven/ui/components/list_view.py b/plugins/ui/src/deephaven/ui/components/list_view.py
index f6985861d..d2434f9dd 100644
--- a/plugins/ui/src/deephaven/ui/components/list_view.py
+++ b/plugins/ui/src/deephaven/ui/components/list_view.py
@@ -7,12 +7,20 @@
 from .item import ItemElement
 from .item_table_source import ItemTableSource
 from ..elements import BaseElement, Element
-from .._internal.utils import create_props
+from .._internal.utils import create_props, unpack_item_table_source
 from ..types import Stringable, Selection, SelectionMode
 
 ListViewItem = Union[Stringable, ItemElement]
 ListViewElement = Element
 
+SUPPORTED_SOURCE_ARGS = {
+    "key_column",
+    "label_column",
+    "description_column",
+    "icon_column",
+    "actions",
+}
+
 
 def list_view(
     *children: ListViewItem | Table | ItemTableSource,
@@ -58,4 +66,6 @@ def list_view(
     """
     children, props = create_props(locals())
 
+    children, props = unpack_item_table_source(children, props, SUPPORTED_SOURCE_ARGS)
+
     return BaseElement("deephaven.ui.components.ListView", *children, **props)
diff --git a/plugins/ui/src/deephaven/ui/components/picker.py b/plugins/ui/src/deephaven/ui/components/picker.py
index 6095df0a9..4ebc5dc3d 100644
--- a/plugins/ui/src/deephaven/ui/components/picker.py
+++ b/plugins/ui/src/deephaven/ui/components/picker.py
@@ -6,11 +6,19 @@
 from .section import SectionElement, PickerItem
 from .item_table_source import ItemTableSource
 from ..elements import BaseElement
-from .._internal.utils import create_props
+from .._internal.utils import create_props, unpack_item_table_source
 from ..types import Key
 
 PickerElement = BaseElement
 
+SUPPORTED_SOURCE_ARGS = {
+    "key_column",
+    "label_column",
+    "description_column",
+    "icon_column",
+    "title_column",
+}
+
 
 def picker(
     *children: PickerItem | SectionElement | Table | PartitionedTable | ItemTableSource,
@@ -53,4 +61,6 @@ def picker(
     """
     children, props = create_props(locals())
 
+    children, props = unpack_item_table_source(children, props, SUPPORTED_SOURCE_ARGS)
+
     return BaseElement("deephaven.ui.components.Picker", *children, **props)
diff --git a/plugins/ui/test/deephaven/ui/test_utils.py b/plugins/ui/test/deephaven/ui/test_utils.py
index f67f2bb80..4d23dd545 100644
--- a/plugins/ui/test/deephaven/ui/test_utils.py
+++ b/plugins/ui/test/deephaven/ui/test_utils.py
@@ -1,4 +1,5 @@
 import unittest
+
 from .BaseTest import BaseTestCase
 
 
@@ -217,6 +218,46 @@ def test_func_with_all_args(a, /, b, *args, c=1, **kwargs):
         # Test that wrapping a function without a signature doesn't throw an error
         wrapped = wrap_callable(print)
 
+    def test_unpack_item_table_source(self):
+        from deephaven.ui._internal.utils import unpack_item_table_source
+        from deephaven.ui import item_table_source
+
+        children = ("table",)
+        props = {
+            "test": "foo",
+        }
+
+        expected_children = ("table",)
+        expected_props = {
+            "test": "foo",
+        }
+
+        self.assertTupleEqual(
+            unpack_item_table_source(children, props, {}),
+            (expected_children, expected_props),
+        )
+
+        item_data_source = item_table_source(
+            table="table", key_column="key", actions="actions"
+        )
+
+        children = (item_data_source,)
+        props = {
+            "test": "foo",
+        }
+
+        expected_children = ("table",)
+
+        expected_props = {
+            "test": "foo",
+            "key_column": "key",
+        }
+
+        self.assertTupleEqual(
+            unpack_item_table_source(children, props, {"table", "key_column"}),
+            (expected_children, expected_props),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()