support doc for multi variants

Author: y1xiaoc

dargs/dargs.py

@@ -335,8 +335,9 @@ def gen_doc_body(self, paths: Optional[List[str]] = None, **kwargs) -> str:
             for subarg in self.sub_fields.values():
                 body_list.append(subarg.gen_doc(paths, **kwargs))
         if self.sub_variants:
+            showflag = len(self.sub_variants) > 1
             for subvrnt in self.sub_variants.values():
-                body_list.append(subvrnt.gen_doc(paths, **kwargs))
+                body_list.append(subvrnt.gen_doc(paths, showflag, **kwargs))
         body = "\n".join(body_list)
         return body
 
@@ -439,13 +440,15 @@ def _convert_choice_alias(self, argdict: dict):
     # above are traversing part
     # below are doc generation part
 
-    def gen_doc(self, paths: Optional[List[str]] = None, **kwargs) -> str:
+    def gen_doc(self, paths: Optional[List[str]] = None, 
+                      showflag : bool = False, **kwargs) -> str:
         body_list = [""]
         body_list.append(f"Depending on the value of *{self.flag_name}*, "
                           "different sub args are accepted. \n") 
         body_list.append(self.gen_doc_flag(paths, **kwargs))
         for choice in self.choice_dict.values():
-            c_str = f"[{choice.name}]"
+            f_str = f"{self.flag_name}=" if showflag else ""
+            c_str = f"[{f_str}{choice.name}]"
             choice_path = [*paths[:-1], paths[-1]+c_str] if paths else [c_str]
             body_list.append("")
             if kwargs.get("make_anchor"):

tests/test_docgen.py

@@ -68,6 +68,50 @@ def test_sub_variants(self):
         docstr = ca.gen_doc(make_anchor=True)
         # print("\n\n"+docstr)
 
+    def test_multi_variants(self):
+        ca = Argument("base", dict, [
+            Argument("sub1", int),
+            Argument("sub2", str)
+        ], [
+            Variant("vnt_flag", [
+                Argument("type1", dict, [
+                    Argument("shared", int),
+                    Argument("vnt1_1", int),
+                    Argument("vnt1_2", dict, [
+                        Argument("vnt1_1_1", int)
+                    ])
+                ]),
+                Argument("type2", dict, [
+                    Argument("shared", int),
+                    Argument("vnt2_1", int),
+                ]),
+                Argument("type3", dict, [
+                    Argument("vnt3_1", int)
+                ], [ # testing cascade variants here
+                    Variant("vnt3_flag1", [
+                        Argument("v3f1t1", dict, [
+                            Argument('v3f1t1_1', int),
+                            Argument('v3f1t1_2', int)
+                        ]),
+                        Argument("v3f1t2", dict, [
+                            Argument('v3f1t2_1', int)
+                        ])
+                    ]),
+                    Variant("vnt3_flag2", [
+                        Argument("v3f2t1", dict, [
+                            Argument('v3f2t1_1', int),
+                            Argument('v3f2t1_2', int)
+                        ]),
+                        Argument("v3f2t2", dict, [
+                            Argument('v3f2t2_1', int)
+                        ])
+                    ])
+                ])
+            ])
+        ])
+        docstr = ca.gen_doc()
+        # print("\n\n"+docstr)
+
     def test_dpmd(self):
         from dpmdargs import gen_doc
         docstr = gen_doc(make_anchor=True)