From 357d91053ff555d615e614ed37cac5f7b403e243 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Wed, 20 Nov 2024 22:25:07 -0500
Subject: [PATCH] fix: fix repeat dict for nb (#86)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **Documentation**
- Updated the Jupyter Notebook to improve the clarity of the JSON
structure by renaming keys and adding new entries for better
organization of repeatable items.
- **Enhancements**
- Improved flexibility of the `ArgumentData` class to handle various
data structures more effectively.
- Refined error handling in the `print_html` method for better
robustness against unknown data types.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 dargs/notebook.py | 24 ++++++++++++++++++++----
 docs/nb.ipynb     |  6 +++++-
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/dargs/notebook.py b/dargs/notebook.py
index c60c366..9d1849f 100644
--- a/dargs/notebook.py
+++ b/dargs/notebook.py
@@ -147,17 +147,24 @@ class ArgumentData:
         The data to be displayed.
     arg : Union[dargs.Argument, dargs.Variant]
         The Argument that describes the data.
+    repeat : bool, optional
+        The argument is repeat
     """
 
-    def __init__(self, data: dict, arg: Argument | Variant):
+    def __init__(self, data: dict, arg: Argument | Variant, repeat: bool = False):
         self.data = data
         self.arg = arg
+        self.repeat = repeat
         self.subdata = []
         self._init_subdata()
 
     def _init_subdata(self):
         """Initialize sub ArgumentData."""
-        if isinstance(self.data, dict) and isinstance(self.arg, Argument):
+        if (
+            isinstance(self.data, dict)
+            and isinstance(self.arg, Argument)
+            and not (self.arg.repeat and not self.repeat)
+        ):
             sub_fields = self.arg.sub_fields.copy()
             # extend subfiles with sub_variants
             for vv in self.arg.sub_variants.values():
@@ -178,9 +185,18 @@ def _init_subdata(self):
             isinstance(self.data, list)
             and isinstance(self.arg, Argument)
             and self.arg.repeat
+            and not self.repeat
         ):
             for dd in self.data:
-                self.subdata.append(ArgumentData(dd, self.arg))
+                self.subdata.append(ArgumentData(dd, self.arg, repeat=True))
+        elif (
+            isinstance(self.data, dict)
+            and isinstance(self.arg, Argument)
+            and self.arg.repeat
+            and not self.repeat
+        ):
+            for dd in self.data.values():
+                self.subdata.append(ArgumentData(dd, self.arg, repeat=True))
 
     def print_html(self, _level=0, _last_one=True):
         """Print the data with Argument in HTML format.
@@ -203,7 +219,7 @@ def print_html(self, _level=0, _last_one=True):
         if _level > 0 and not (
             isinstance(self.data, dict)
             and isinstance(self.arg, Argument)
-            and self.arg.repeat
+            and self.repeat
         ):
             if isinstance(self.arg, (Argument, Variant)):
                 buff.append(r"""<span class="dargs-key">""")
diff --git a/docs/nb.ipynb b/docs/nb.ipynb
index 9350ffd..816d94f 100644
--- a/docs/nb.ipynb
+++ b/docs/nb.ipynb
@@ -29,10 +29,14 @@
     "        3\n",
     "    ],\n",
     "    \"test_variant\": \"test_variant_argument\",\n",
-    "    \"test_repeat\": [\n",
+    "    \"test_repeat_list\": [\n",
     "        {\"test_repeat_item\": false},\n",
     "        {\"test_repeat_item\": true}\n",
     "    ],\n",
+    "    \"test_repeat_dict\": {\n",
+    "        \"test1\": {\"test_repeat_item\": false},\n",
+    "        \"test2\": {\"test_repeat_item\": true}\n",
+    "    },\n",
     "    \"_comment\": \"This is an example data\"\n",
     "}\n",
     "\"\"\"\n",