Enhancing Choice/Case options

This patch allows user to use with_choices option in children() calls
CESNET · Jan 26, 2024 · 3dd971e · 3dd971e
1 parent 3e3af68
commit 3dd971e
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 12 deletions.
diff --git a/cffi/cdefs.h b/cffi/cdefs.h
@@ -1006,6 +1006,18 @@ struct lysc_when {
 
 struct lysc_when** lysc_node_when(const struct lysc_node *);
 
+struct lysc_node_case {
+    struct lysc_node *child;
+    struct lysc_when **when;
+    ...;
+};
+
+struct lysc_node_choice {
+    struct lysc_node_case *cases;
+    struct lysc_when **when;
+    ...;
+};
+
 #define LYD_DEFAULT ...
 #define LYD_WHEN_TRUE ...
 #define LYD_NEW ...

diff --git a/libyang/schema.py b/libyang/schema.py
@@ -137,8 +137,11 @@ 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_choices: bool = False
+    ) -> Iterator["SNode"]:
+        options = lib.LYS_GETNEXT_WITHCHOICE if with_choices else 0
+        return iter_children(self.context, self.cdata, types=types, options=options)
 
     def __str__(self) -> str:
         return self.name()
@@ -1357,18 +1360,30 @@ 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_choices: bool = False
+    ) -> Iterator[SNode]:
+        options = lib.LYS_GETNEXT_WITHCHOICE if with_choices else 0
+        return iter_children(self.context, self.cdata, types=types, options=options)
 
 
 # -------------------------------------------------------------------------------------
 @SNode.register(SNode.CHOICE)
 class SChoice(SNode):
+    __slots__ = ("cdata_choice",)
+
+    def __init__(self, context: "libyang.Context", cdata):
+        super().__init__(context, cdata)
+        self.cdata_choice = ffi.cast("struct lysc_node_choice *", cdata)
+
     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]:
+        options = lib.LYS_GETNEXT_WITHCASE if with_cases else 0
+        return iter_children(self.context, self.cdata, types=types, options=options)
 
 
 # -------------------------------------------------------------------------------------
@@ -1377,8 +1392,11 @@ 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_choices: bool = False
+    ) -> Iterator[SNode]:
+        options = lib.LYS_GETNEXT_WITHCHOICE if with_choices else 0
+        return iter_children(self.context, self.cdata, types=types, options=options)
 
 
 # -------------------------------------------------------------------------------------
@@ -1398,9 +1416,15 @@ 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_choices: bool = False,
     ) -> Iterator[SNode]:
-        return iter_children(self.context, self.cdata, skip_keys=skip_keys, types=types)
+        options = lib.LYS_GETNEXT_WITHCHOICE if with_choices else 0
+        return iter_children(
+            self.context, self.cdata, skip_keys=skip_keys, types=types, options=options
+        )
 
     def keys(self) -> Iterator[SNode]:
         node = lib.lysc_node_child(self.cdata)
@@ -1498,15 +1522,19 @@ def iter_children(
     options: int = 0,
 ) -> Iterator[SNode]:
     if types is None:
-        types = (
+        types = [
             lib.LYS_ACTION,
             lib.LYS_CONTAINER,
             lib.LYS_LIST,
             lib.LYS_RPC,
             lib.LYS_LEAF,
             lib.LYS_LEAFLIST,
             lib.LYS_NOTIF,
-        )
+        ]
+        if options & lib.LYS_GETNEXT_WITHCHOICE:
+            types.append(lib.LYS_CHOICE)
+        if options & lib.LYS_GETNEXT_WITHCASE:
+            types.append(lib.LYS_CASE)
 
     def _skip(node) -> bool:
         if node.nodetype not in types:

diff --git a/tests/test_schema.py b/tests/test_schema.py
@@ -269,6 +269,10 @@ 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())