Enhancing Choice/Case options

This patch allows user to use with_choices option in children() calls

libyang/schema.py

@@ -137,8 +137,12 @@ def get_module_from_prefix(self, prefix: str) -> Optional["Module"]:
     def __iter__(self) -> Iterator["SNode"]:
         return self.children()
 
-    def children(self, types: Optional[Tuple[int, ...]] = None) -> Iterator["SNode"]:
-        return iter_children(self.context, self.cdata, types=types)
+    def children(
+        self, types: Optional[Tuple[int, ...]] = None, with_choice: bool = False
+    ) -> Iterator["SNode"]:
+        return iter_children(
+            self.context, self.cdata, types=types, with_choice=with_choice
+        )
 
     def __str__(self) -> str:
         return self.name()
@@ -1357,8 +1361,12 @@ def must_conditions(self) -> Iterator[str]:
     def __iter__(self) -> Iterator[SNode]:
         return self.children()
 
-    def children(self, types: Optional[Tuple[int, ...]] = None) -> Iterator[SNode]:
-        return iter_children(self.context, self.cdata, types=types)
+    def children(
+        self, types: Optional[Tuple[int, ...]] = None, with_choice: bool = False
+    ) -> Iterator[SNode]:
+        return iter_children(
+            self.context, self.cdata, types=types, with_choice=with_choice
+        )
 
 
 # -------------------------------------------------------------------------------------
@@ -1367,8 +1375,10 @@ class SChoice(SNode):
     def __iter__(self) -> Iterator[SNode]:
         return self.children()
 
-    def children(self, types: Optional[Tuple[int, ...]] = None) -> Iterator[SNode]:
-        return iter_children(self.context, self.cdata, types=types)
+    def children(
+        self, types: Optional[Tuple[int, ...]] = None, with_cases: bool = False
+    ) -> Iterator[SNode]:
+        return iter_children(self.context, self.cdata, types=types, with_case=with_case)
 
 
 # -------------------------------------------------------------------------------------
@@ -1377,8 +1387,12 @@ class SCase(SNode):
     def __iter__(self) -> Iterator[SNode]:
         return self.children()
 
-    def children(self, types: Optional[Tuple[int, ...]] = None) -> Iterator[SNode]:
-        return iter_children(self.context, self.cdata, types=types)
+    def children(
+        self, types: Optional[Tuple[int, ...]] = None, with_choice: bool = False
+    ) -> Iterator[SNode]:
+        return iter_children(
+            self.context, self.cdata, types=types, with_choice=with_choice
+        )
 
 
 # -------------------------------------------------------------------------------------
@@ -1398,9 +1412,18 @@ def __iter__(self) -> Iterator[SNode]:
         return self.children()
 
     def children(
-        self, skip_keys: bool = False, types: Optional[Tuple[int, ...]] = None
+        self,
+        skip_keys: bool = False,
+        types: Optional[Tuple[int, ...]] = None,
+        with_choice: bool = False,
     ) -> Iterator[SNode]:
-        return iter_children(self.context, self.cdata, skip_keys=skip_keys, types=types)
+        return iter_children(
+            self.context,
+            self.cdata,
+            skip_keys=skip_keys,
+            types=types,
+            with_choice=with_choice,
+        )
 
     def keys(self) -> Iterator[SNode]:
         node = lib.lysc_node_child(self.cdata)
@@ -1489,13 +1512,43 @@ class SAnydata(SNode):
     pass
 
 
+# -------------------------------------------------------------------------------------
+def iter_children_options(
+    with_choice: bool = False,
+    no_choice: bool = False,
+    with_case: bool = False,
+    into_non_presence_container: bool = False,
+    output: bool = False,
+    with_schema_mount: bool = False,
+) -> int:
+    options = 0
+    if with_choice:
+        options |= lib.LYS_GETNEXT_WITHCHOICE
+    if no_choice:
+        options |= lib.LYS_GETNEXT_NOCHOICE
+    if with_case:
+        options |= lib.LYS_GETNEXT_WITHCASE
+    if into_non_presence_container:
+        options |= lib.LYS_GETNEXT_INTONPCONT
+    if output:
+        options |= lib.LYS_GETNEXT_OUTPUT
+    if with_schema_mount:
+        options |= lib.LYS_GETNEXT_WITHSCHEMAMOUNT
+    return options
+
+
 # -------------------------------------------------------------------------------------
 def iter_children(
     context: "libyang.Context",
     parent,  # C type: Union["struct lys_module *", "struct lys_node *"]
     skip_keys: bool = False,
     types: Optional[Tuple[int, ...]] = None,
-    options: int = 0,
+    with_choice: bool = False,
+    no_choice: bool = False,
+    with_case: bool = False,
+    into_non_presence_container: bool = False,
+    output: bool = False,
+    with_schema_mount: bool = False,
 ) -> Iterator[SNode]:
     if types is None:
         types = (
@@ -1506,6 +1559,8 @@ def iter_children(
             lib.LYS_LEAF,
             lib.LYS_LEAFLIST,
             lib.LYS_NOTIF,
+            lib.LYS_CHOICE,
+            lib.LYS_CASE,
         )
 
     def _skip(node) -> bool:
@@ -1526,6 +1581,14 @@ def _skip(node) -> bool:
     else:
         module = ffi.NULL
 
+    options = iter_children_options(
+        with_choice=with_choice,
+        no_choice=no_choice,
+        with_case=with_case,
+        into_non_presence_container=into_non_presence_container,
+        output=output,
+        with_schema_mount=with_schema_mount,
+    )
     child = lib.lys_getnext(ffi.NULL, parent, module, options)
     while child:
         if not _skip(child):

tests/test_schema.py

@@ -269,6 +269,12 @@ def test_cont_iter(self):
     def test_cont_children_leafs(self):
         leafs = list(self.container.children(types=(SNode.LEAF,)))
         self.assertEqual(len(leafs), 9)
+        without_choices = [
+            c.name() for c in self.container.children(with_choices=False)
+        ]
+        with_choices = [c.name() for c in self.container.children(with_choices=True)]
+        self.assertTrue("pill" not in without_choices)
+        self.assertTrue("pill" in with_choices)
 
     def test_cont_parent(self):
         self.assertIsNone(self.container.parent())