feat: add update function to JSONPath class (#12)

Author: sean2077

jsonpath/jsonpath.py

@@ -3,7 +3,7 @@
 import re
 import sys
 from collections import defaultdict
-from typing import Union
+from typing import Any, Callable, Union
 
 
 def create_logger(name: str = None, level: Union[int, str] = logging.INFO):
@@ -59,6 +59,7 @@ class JSONPath:
     REP_SLICE_CONTENT = re.compile(r"^(-?\d*)?:(-?\d*)?(:-?\d*)?$")
     REP_SELECT_CONTENT = re.compile(r"^([\w.']+)(, ?[\w.']+)+$")
     REP_FILTER_CONTENT = re.compile(r"@([.\[].*?)(?=<=|>=|==|!=|>|<| in| not| is|\s|\)|$)|len\(@([.\[].*?)\)")
+    REP_PATH_SEGMENT = re.compile(r"(?:\.|^)(?P<dot>\w+)|\[['\"](?P<quote>.*?)['\"]\]|\[(?P<int>\d+)\]")
 
     # annotations
     f: list
@@ -110,7 +111,9 @@ def _parse_expr(self, expr):
         expr = JSONPath.REP_PUT_PAREN.sub(self._put_paren, expr)
         expr = JSONPath.REP_PUT_BACKQUOTE.sub(self._put_backquote, expr)
         expr = JSONPath.REP_PUT_QUOTE.sub(self._put_quote, expr)
-        if expr.startswith("$;"):
+        if expr == "$":
+            expr = ""
+        elif expr.startswith("$;"):
             expr = expr[2:]
 
         logger.debug(f"after expr  : {expr}")
@@ -335,6 +338,46 @@ def _trace(self, obj, i: int, path):
 
             return
 
+    def update(self, obj: Union[list, dict], value_or_func: Union[Any, Callable[[Any], Any]]) -> Any:
+        paths = self.parse(obj, result_type="PATH")
+        for path in paths:
+            matches = list(JSONPath.REP_PATH_SEGMENT.finditer(path))
+            if not matches:
+                # Root object
+                if isinstance(value_or_func, Callable):
+                    obj = value_or_func(obj)
+                else:
+                    obj = value_or_func
+                continue
+
+            target = obj
+            # Traverse to parent
+            for match in matches[:-1]:
+                group = match.groupdict()
+                if group["dot"]:
+                    target = target[group["dot"]]
+                elif group["quote"]:
+                    target = target[group["quote"]]
+                elif group["int"]:
+                    target = target[int(group["int"])]
+
+            # Update last segment
+            last_match = matches[-1]
+            group = last_match.groupdict()
+            if group["dot"]:
+                key = group["dot"]
+            elif group["quote"]:
+                key = group["quote"]
+            elif group["int"]:
+                key = int(group["int"])
+
+            if isinstance(value_or_func, Callable):
+                target[key] = value_or_func(target[key])
+            else:
+                target[key] = value_or_func
+
+        return obj
+
 
 class RegexPattern:
     def __init__(self, pattern):

tests/test_update.py

@@ -0,0 +1,126 @@
+from jsonpath.jsonpath import JSONPath
+
+
+def test_update_value():
+    data = {
+        "store": {
+            "book": [
+                {"category": "reference", "author": "Nigel Rees", "title": "Sayings of the Century", "price": 8.95}
+            ]
+        }
+    }
+    jp = JSONPath("$.store.book[0].price")
+    result = jp.update(data, 10.0)
+    assert result["store"]["book"][0]["price"] == 10.0
+
+
+def test_update_function():
+    data = {"count": 1}
+    jp = JSONPath("$.count")
+    result = jp.update(data, lambda x: x + 1)
+    assert result["count"] == 2
+
+
+def test_update_root():
+    data = {"a": 1}
+    jp = JSONPath("$")
+    result = jp.update(data, {"b": 2})
+    assert result == {"b": 2}
+
+
+def test_update_list_index():
+    data = [1, 2, 3]
+    jp = JSONPath("$[1]")
+    result = jp.update(data, 5)
+    assert result == [1, 5, 3]
+
+
+def test_update_multiple():
+    data = {"items": [{"value": 1}, {"value": 2}, {"value": 3}]}
+    jp = JSONPath("$.items[*].value")
+    result = jp.update(data, 0)
+    for item in result["items"]:
+        assert item["value"] == 0
+
+
+def test_update_multiple_func():
+    data = {"items": [{"value": 1}, {"value": 2}, {"value": 3}]}
+    jp = JSONPath("$.items[*].value")
+    result = jp.update(data, lambda x: x * 2)
+    assert result["items"][0]["value"] == 2
+    assert result["items"][1]["value"] == 4
+    assert result["items"][2]["value"] == 6
+
+
+def test_update_with_filter():
+    data = {"books": [{"price": 10, "title": "A"}, {"price": 20, "title": "B"}, {"price": 30, "title": "C"}]}
+    # Update price where price > 15
+    jp = JSONPath("$.books[?(@.price > 15)].price")
+    result = jp.update(data, 0)
+    assert result["books"][0]["price"] == 10  # Unchanged
+    assert result["books"][1]["price"] == 0  # Updated
+    assert result["books"][2]["price"] == 0  # Updated
+
+
+def test_update_slice():
+    data = [0, 1, 2, 3, 4]
+    jp = JSONPath("$[1:4]")  # Indices 1, 2, 3
+    result = jp.update(data, 9)
+    assert result == [0, 9, 9, 9, 4]
+
+
+def test_update_special_keys():
+    data = {"complex.key": 1, "key with space": 2, "normal": 3}
+    # Note: jsonpath-python might handle keys with dots using ['...'] syntax in path output
+    jp = JSONPath("$['complex.key']")
+    result = jp.update(data, 10)
+    assert result["complex.key"] == 10
+
+    jp = JSONPath("$['key with space']")
+    result = jp.update(data, 20)
+    assert result["key with space"] == 20
+
+
+def test_update_no_match():
+    data = {"a": 1}
+    jp = JSONPath("$.b")
+    result = jp.update(data, 2)
+    assert result == {"a": 1}
+
+
+def test_update_nested_structure():
+    data = {"a": [{"b": [1, 2]}, {"b": [3, 4]}]}
+    jp = JSONPath("$.a[*].b[1]")
+    result = jp.update(data, 99)
+    assert result["a"][0]["b"][1] == 99
+    assert result["a"][1]["b"][1] == 99
+
+
+def test_update_recursive():
+    data = {
+        "store": {
+            "book": [{"category": "reference", "price": 8.95}, {"category": "fiction", "price": 12.99}],
+            "bicycle": {"color": "red", "price": 19.95},
+        }
+    }
+    jp = JSONPath("$..price")
+    result = jp.update(data, 10.0)
+    assert result["store"]["book"][0]["price"] == 10.0
+    assert result["store"]["book"][1]["price"] == 10.0
+    assert result["store"]["bicycle"]["price"] == 10.0
+
+
+def test_update_union():
+    data = {"a": 1, "b": 2, "c": 3}
+    jp = JSONPath("$[a,b]")
+    result = jp.update(data, 0)
+    assert result["a"] == 0
+    assert result["b"] == 0
+    assert result["c"] == 3
+
+
+def test_update_quote_keys():
+    data = {'c"d': 1, "normal": 2}
+    jp = JSONPath("$['c\"d']")
+    result = jp.update(data, 99)
+    assert result['c"d'] == 99