Add simple query projection

Author: jg-rp

jsonpath/env.py

@@ -353,7 +353,7 @@ def query(
             ...
         ```
         """
-        return Query(self.finditer(path, data, filter_context=filter_context))
+        return Query(self.finditer(path, data, filter_context=filter_context), self)
 
     async def findall_async(
         self,

jsonpath/fluent_api.py

@@ -5,12 +5,22 @@
 import collections
 import itertools
 from typing import TYPE_CHECKING
+from typing import Any
+from typing import Dict
 from typing import Iterable
 from typing import Iterator
+from typing import List
+from typing import Mapping
 from typing import Optional
+from typing import Sequence
 from typing import Tuple
+from typing import Union
+
+from .exceptions import JSONPointerKeyError
+from .patch import JSONPatch
 
 if TYPE_CHECKING:
+    from jsonpath import JSONPathEnvironment
     from jsonpath import JSONPathMatch
     from jsonpath import JSONPointer
 
@@ -28,8 +38,9 @@ class Query:
     **New in version 1.1.0**
     """
 
-    def __init__(self, it: Iterable[JSONPathMatch]) -> None:
+    def __init__(self, it: Iterable[JSONPathMatch], env: JSONPathEnvironment) -> None:
         self._it = iter(it)
+        self._env = env
 
     def __iter__(self) -> Iterator[JSONPathMatch]:
         return self._it
@@ -126,6 +137,31 @@ def pointers(self) -> Iterable[JSONPointer]:
         """Return an iterable of JSONPointers, one for each match."""
         return (m.pointer() for m in self._it)
 
+    def select(self, *expressions: str) -> Iterable[object]:
+        """Query projection using relative JSONPaths.
+
+        Returns an iterable of objects built from selecting _expressions_ relative to
+        each match from the current query.
+        """
+        for m in self._it:
+            if isinstance(m.obj, Sequence):
+                obj: Union[List[Any], Dict[str, Any]] = []
+            elif isinstance(m.obj, Mapping):
+                obj = {}
+            else:
+                return iter([])
+
+            patch = JSONPatch()
+
+            for expr in expressions:
+                for match in self._env.finditer(expr, m.obj):  # type: ignore
+                    _pointer = match.pointer()
+                    _patch_parents(_pointer.parent(), patch, m.obj)  # type: ignore
+                    patch.add(_pointer, match.obj)
+
+            patch.apply(obj)
+            yield obj
+
     def first_one(self) -> Optional[JSONPathMatch]:
         """Return the first `JSONPathMatch` or `None` if there were no matches."""
         try:
@@ -152,11 +188,31 @@ def tee(self, n: int = 2) -> Tuple[Query, ...]:
 
         It is not safe to use a `Query` instance after calling `tee()`.
         """
-        return tuple(Query(it) for it in itertools.tee(self._it, n))
+        return tuple(Query(it, self._env) for it in itertools.tee(self._it, n))
 
     def take(self, n: int) -> Query:
         """Return a new query iterating over the next _n_ matches.
 
         It is safe to continue using this query after calling take.
         """
-        return Query(list(itertools.islice(self._it, n)))
+        return Query(list(itertools.islice(self._it, n)), self._env)
+
+
+def _patch_parents(
+    pointer: JSONPointer,
+    patch: JSONPatch,
+    obj: Union[Sequence[Any], Mapping[str, Any]],
+) -> None:
+    if pointer.parent().parts:
+        _patch_parents(pointer.parent(), patch, obj)
+
+    try:
+        _obj = pointer.resolve(obj)
+    except JSONPointerKeyError:
+        _obj = obj
+
+    if pointer.parts:
+        if isinstance(_obj, Sequence):
+            patch.addne(pointer, [])
+        elif isinstance(_obj, Mapping):
+            patch.addne(pointer, {})

jsonpath/patch.py

@@ -85,6 +85,52 @@ def asdict(self) -> Dict[str, object]:
         return {"op": self.name, "path": str(self.path), "value": self.value}
 
 
+class OpAddNe(Op):
+    """A non-standard _add if not exists_ operation."""
+
+    __slots__ = ("path", "value")
+
+    name = "add"
+
+    def __init__(self, path: JSONPointer, value: object) -> None:
+        self.path = path
+        self.value = value
+
+    def apply(
+        self, data: Union[MutableSequence[object], MutableMapping[str, object]]
+    ) -> Union[MutableSequence[object], MutableMapping[str, object]]:
+        """Apply this patch operation to _data_."""
+        parent, obj = self.path.resolve_parent(data)
+        if parent is None:
+            # Replace the root object.
+            # The following op, if any, will raise a JSONPatchError if needed.
+            return self.value  # type: ignore
+
+        target = self.path.parts[-1]
+        if isinstance(parent, MutableSequence):
+            if obj is UNDEFINED:
+                if target == "-":
+                    parent.append(self.value)
+                else:
+                    raise JSONPatchError("index out of range")
+            else:
+                parent.insert(int(target), self.value)
+        elif (
+            isinstance(parent, MutableMapping)
+            and parent.get(target, UNDEFINED) == UNDEFINED
+        ):
+            parent[target] = self.value
+        # else:
+        #     raise JSONPatchError(
+        #         f"unexpected operation on {parent.__class__.__name__!r}"
+        #     )
+        return data
+
+    def asdict(self) -> Dict[str, object]:
+        """Return a dictionary representation of this operation."""
+        return {"op": self.name, "path": str(self.path), "value": self.value}
+
+
 class OpRemove(Op):
     """The JSON Patch _remove_ operation."""
 
@@ -340,6 +386,11 @@ def _build(self, patch: Iterable[Mapping[str, object]]) -> None:
                     path=self._op_pointer(operation, "path", "add", i),
                     value=self._op_value(operation, "value", "add", i),
                 )
+            elif op == "addne":
+                self.addne(
+                    path=self._op_pointer(operation, "path", "add", i),
+                    value=self._op_value(operation, "value", "add", i),
+                )
             elif op == "remove":
                 self.remove(path=self._op_pointer(operation, "path", "add", i))
             elif op == "replace":
@@ -424,6 +475,22 @@ def add(self: Self, path: Union[str, JSONPointer], value: object) -> Self:
         self.ops.append(OpAdd(path=pointer, value=value))
         return self
 
+    def addne(self: Self, path: Union[str, JSONPointer], value: object) -> Self:
+        """Append an _addne_ operation to this patch.
+
+        Arguments:
+            path: A string representation of a JSON Pointer, or one that has
+                already been parsed.
+            value: The object to add.
+
+        Returns:
+            This `JSONPatch` instance, so we can build a JSON Patch by chaining
+                calls to JSON Patch operation methods.
+        """
+        pointer = self._ensure_pointer(path)
+        self.ops.append(OpAddNe(path=pointer, value=value))
+        return self
+
     def remove(self: Self, path: Union[str, JSONPointer]) -> Self:
         """Append a _remove_ operation to this patch.