xpath: add support for indexed xpath

This patch adds support for indexed xpath to xpath_set.
In the case of a leaf-list, it is now possible to apply xpath_set in this way:

xpath_set(d, "/lstnum[5]", 33, after="4")

This will replace the 5th element of the leaf-list with the new value 33.

This feature is important because libyang can return this type of xpath
when a callback is called with sr_module_change_subscribe from sysrepo.

The tests are updated accordingly.

Signed-off-by: Jean-Sébastien Bevilacqua <jean-sebastien.bevilacqua@6wind.com>

Author: jeanseb6wind

libyang/xpath.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2020 6WIND S.A.
 # SPDX-License-Identifier: MIT
 
+import contextlib
 import fnmatch
 import re
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
@@ -56,17 +57,28 @@ def xpath_split(xpath: str) -> Iterator[Tuple[str, str, List[Tuple[str, str]]]]:
         while i < len(xpath) and xpath[i] == "[":
             i += 1  # skip opening '['
             j = xpath.find("=", i)  # find key name end
-            key_name = xpath[i:j]
-            quote = xpath[j + 1]  # record opening quote character
-            j = i = j + 2  # skip '=' and opening quote
-            while True:
-                if xpath[j] == quote and xpath[j - 1] != "\\":
-                    break
-                j += 1
-            # replace escaped chars by their non-escape version
-            key_value = xpath[i:j].replace(f"\\{quote}", f"{quote}")
-            keys.append((key_name, key_value))
-            i = j + 2  # skip closing quote and ']'
+
+            if j != -1:  # keyed specifier
+                key_name = xpath[i:j]
+                quote = xpath[j + 1]  # record opening quote character
+                j = i = j + 2  # skip '=' and opening quote
+                while True:
+                    if xpath[j] == quote and xpath[j - 1] != "\\":
+                        break
+                    j += 1
+                # replace escaped chars by their non-escape version
+                key_value = xpath[i:j].replace(f"\\{quote}", f"{quote}")
+                keys.append((key_name, key_value))
+                i = j + 2  # skip closing quote and ']'
+            else:  # index specifier
+                j = i
+                while True:
+                    if xpath[j] == "]":
+                        break
+                    j += 1
+                key_value = xpath[i:j]
+                keys.append(("", key_value))
+                i = j + 2
 
         yield prefix, name, keys
 
@@ -134,6 +146,12 @@ def _list_find_key_index(keys: List[Tuple[str, str]], lst: List) -> int:
             if py_to_yang(elem) == keys[0][1]:
                 return i
 
+    elif keys[0][0] == "":
+        # keys[0][1] is directly the index
+        index = int(keys[0][1]) - 1
+        if len(lst) > index:
+            return index
+
     else:
         for i, elem in enumerate(lst):
             if not isinstance(elem, dict):
@@ -410,32 +428,47 @@ def xpath_set(
             lst.append(value)
         return lst[key_val]
 
-    if isinstance(lst, list):
-        # regular python list, need to iterate over it
-        try:
-            i = _list_find_key_index(keys, lst)
-            # found
-            if force:
-                lst[i] = value
-            return lst[i]
-        except ValueError:
-            # not found
-            if after is None:
-                lst.append(value)
-            elif after == "":
-                lst.insert(0, value)
-            else:
-                if after[0] != "[":
-                    after = "[.=%r]" % str(after)
-                _, _, after_keys = next(xpath_split("/*" + after))
-                insert_index = _list_find_key_index(after_keys, lst) + 1
-                if insert_index == len(lst):
-                    lst.append(value)
-                else:
-                    lst.insert(insert_index, value)
-            return value
+    # regular python list from now
+    if not isinstance(lst, list):
+        raise TypeError("expected a list")
+
+    with contextlib.suppress(ValueError):
+        i = _list_find_key_index(keys, lst)
+        # found
+        if force:
+            lst[i] = value
+        return lst[i]
+
+    # value not found; handle insertion based on 'after'
+    if after is None:
+        lst.append(value)
+        return value
+
+    if after == "":
+        lst.insert(0, value)
+        return value
+
+    # first try to find the value in the leaf list
+    try:
+        _, _, after_keys = next(
+            xpath_split(f"/*{after}" if after[0] == "[" else f"/*[.={after!r}]")
+        )
+        insert_index = _list_find_key_index(after_keys, lst) + 1
+    except ValueError:
+        # handle 'after' as numeric index
+        if not after.isnumeric():
+            raise
+
+        insert_index = int(after)
+        if insert_index > len(lst):
+            raise
+
+    if insert_index == len(lst):
+        lst.append(value)
+    else:
+        lst.insert(insert_index, value)
 
-    raise TypeError("expected a list")
+    return value
 
 
 # -------------------------------------------------------------------------------------

tests/test_xpath.py

@@ -41,6 +41,11 @@ def test_xpath_set(self):
         )
         ly.xpath_set(d, "/lstnum[.='100']", 100)
         ly.xpath_set(d, "/lstnum[.='1']", 1, after="")
+        ly.xpath_set(d, "/lstnum[5]", 33, after="4")
+        ly.xpath_set(d, "/lstnum[5]", 34, after="4")
+        ly.xpath_set(d, "/lstnum[5]", 35, after="4")
+        ly.xpath_set(d, "/lstnum[7]", 101, after="6")
+        ly.xpath_set(d, "/lstnum[8]", 102, after="7")
         with self.assertRaises(ValueError):
             ly.xpath_set(d, "/lstnum[.='1000']", 1000, after="1000000")
         with self.assertRaises(ValueError):
@@ -101,7 +106,7 @@ def test_xpath_set(self):
                     {"name": "eth3", "mtu": 1000},
                 ],
                 "lst2": ["a", "b", "c"],
-                "lstnum": [1, 10, 20, 30, 40, 100],
+                "lstnum": [1, 10, 20, 30, 35, 100, 101, 102],
                 "val": 43,
             },
         )