Choice of projection types

jg-rp · jg-rp · commit c28ff63ddbce · 2024-06-30T15:11:33.000+01:00
diff --git a/jsonpath/__init__.py b/jsonpath/__init__.py
@@ -17,6 +17,7 @@
 from .exceptions import RelativeJSONPointerIndexError
 from .exceptions import RelativeJSONPointerSyntaxError
 from .filter import UNDEFINED
+from .fluent_api import Projection
 from .fluent_api import Query
 from .lex import Lexer
 from .match import JSONPathMatch
@@ -53,6 +54,7 @@
     "Lexer",
     "match",
     "Parser",
+    "Projection",
     "query",
     "Query",
     "RelativeJSONPointer",
diff --git a/jsonpath/fluent_api.py b/jsonpath/fluent_api.py
@@ -4,6 +4,8 @@
 
 import collections
 import itertools
+from enum import Enum
+from enum import auto
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Dict
@@ -25,6 +27,14 @@
     from jsonpath import JSONPointer
 
 
+class Projection(Enum):
+    """Projection style."""
+
+    RELATIVE = auto()
+    ROOT = auto()
+    FLAT = auto()
+
+
 class Query:
     """A fluent API for managing `JSONPathMatch` iterators.
 
@@ -137,27 +147,45 @@ 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]:
+    def select(
+        self,
+        *expressions: str,
+        projection: Projection = Projection.RELATIVE,
+    ) -> 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):
+            if isinstance(m.obj, Sequence) or projection == Projection.FLAT:
                 obj: Union[List[Any], Dict[str, Any]] = []
             elif isinstance(m.obj, Mapping):
                 obj = {}
             else:
                 return
 
+            root_pointer = m.pointer()
             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.addap(_pointer, match.obj)
+                    if projection == Projection.FLAT:
+                        patch.addap("/-", match.obj)
+                    elif projection == Projection.ROOT:
+                        # Pointer string without a leading slash
+                        rel_pointer = "/".join(
+                            str(p).replace("~", "~0").replace("/", "~1")
+                            for p in match.parts
+                        )
+                        _pointer = root_pointer / rel_pointer
+                        _patch_parents(_pointer.parent(), patch, m.obj)  # type: ignore
+                        patch.addap(_pointer, match.obj)
+                    else:
+                        # Natural projection
+                        _pointer = match.pointer()
+                        _patch_parents(_pointer.parent(), patch, m.obj)  # type: ignore
+                        patch.addap(_pointer, match.obj)
 
             patch.apply(obj)
 
diff --git a/tests/test_query_project.py b/tests/test_query_project.py
@@ -4,6 +4,15 @@
 from typing import List
 
 import jsonpath
+from jsonpath import Projection
+
+
+def test_select_from_primitive() -> None:
+    expr = "$[0].a"
+    data = [{"a": 1, "b": 1}, {"a": 2, "b": 2}, {"b": 3, "a": 3}]
+    projection = ("nosuchthing",)
+    it = jsonpath.query(expr, data).select(*projection)
+    assert list(it) == []
 
 
 def test_top_level_array() -> None:
@@ -14,6 +23,14 @@ def test_top_level_array() -> None:
     assert list(it) == [{"a": 1}, {"a": 2}, {"a": 3}]
 
 
+def test_top_level_array_flat_projection() -> None:
+    expr = "$.*"
+    data = [{"a": 1, "b": 1}, {"a": 2, "b": 2}, {"b": 3, "a": 3}]
+    projection = ("a",)
+    it = jsonpath.query(expr, data).select(*projection, projection=Projection.FLAT)
+    assert list(it) == [[1], [2], [3]]
+
+
 def test_top_level_array_partial_existence() -> None:
     expr = "$.*"
     data = [{"a": 1, "b": 1}, {"b": 2}, {"b": 3, "a": 3}]
@@ -76,3 +93,19 @@ def test_select_nested_objects() -> None:
     projection = ("foo.bar",)
     it = jsonpath.query(expr, data).select(*projection)
     assert list(it) == [{"foo": {"bar": 42}}]
+
+
+def test_select_nested_objects_root_projection() -> None:
+    expr = "$.a"
+    data = {"a": {"foo": {"bar": 42}, "bar": 7}, "b": 1}
+    projection = ("foo.bar",)
+    it = jsonpath.query(expr, data).select(*projection, projection=Projection.ROOT)
+    assert list(it) == [{"a": {"foo": {"bar": 42}}}]
+
+
+def test_select_nested_objects_flat_projection() -> None:
+    expr = "$.a"
+    data = {"a": {"foo": {"bar": 42}, "bar": 7}, "b": 1}
+    projection = ("foo.bar",)
+    it = jsonpath.query(expr, data).select(*projection, projection=Projection.FLAT)
+    assert list(it) == [[42]]
