From 586fbc2d54a7ad6bbf543e54a66d558ab4510269 Mon Sep 17 00:00:00 2001
From: James Prior <jamesgr.prior@gmail.com>
Date: Thu, 11 Jul 2024 07:14:01 +0100
Subject: [PATCH] Allow `Query.select()` to accept pre-compiled paths

---
 docs/query.md                  |  8 ++++++--
 jsonpath/fluent_api.py         | 13 ++++++++-----
 tests/test_query_projection.py |  8 ++++++++
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/docs/query.md b/docs/query.md
index c008a70..b49192a 100644
--- a/docs/query.md
+++ b/docs/query.md
@@ -166,12 +166,16 @@ for product in query("$..products.*", data).values():
 {'title': 'Beanie', 'description': 'Winter running hat.', 'price': 9.0}
 ```
 
-We can select nested values too.
+We can select nested values too, and arguments to `select()` can be pre-compiled paths.
 
 ```python
+import jsonpath
+
 # ...
 
-for product in query("$..products.*", data).select("title", "social.shares"):
+projection = (jsonpath.compile("title"), jsonpath.compile("social.shares"))
+
+for product in jsonpath.query("$..products.*", data).select(*projection):
     print(product)
 ```
 
diff --git a/jsonpath/fluent_api.py b/jsonpath/fluent_api.py
index 52a5204..67a4539 100644
--- a/jsonpath/fluent_api.py
+++ b/jsonpath/fluent_api.py
@@ -22,6 +22,8 @@
 from .patch import JSONPatch
 
 if TYPE_CHECKING:
+    from jsonpath import CompoundJSONPath
+    from jsonpath import JSONPath
     from jsonpath import JSONPathEnvironment
     from jsonpath import JSONPathMatch
     from jsonpath import JSONPointer
@@ -191,7 +193,7 @@ def take(self, n: int) -> Query:
 
     def select(
         self,
-        *expressions: str,
+        *expressions: Union[str, JSONPath, CompoundJSONPath],
         projection: Projection = Projection.RELATIVE,
     ) -> Iterable[object]:
         """Query projection using relative JSONPaths.
@@ -217,7 +219,7 @@ def select(
     def _select(
         self,
         match: JSONPathMatch,
-        expressions: Tuple[str, ...],
+        expressions: Tuple[Union[str, JSONPath, CompoundJSONPath], ...],
         projection: Projection,
     ) -> object:
         if isinstance(match.obj, str):
@@ -232,20 +234,21 @@ def _select(
         patch = JSONPatch()
 
         for expr in expressions:
-            self._patch(match, expr, patch, projection)
+            path = self._env.compile(expr) if isinstance(expr, str) else expr
+            self._patch(match, path, patch, projection)
 
         return _fix_sparse_arrays(patch.apply(obj))
 
     def _patch(
         self,
         match: JSONPathMatch,
-        expr: str,
+        path: Union[JSONPath, CompoundJSONPath],
         patch: JSONPatch,
         projection: Projection,
     ) -> None:
         root_pointer = match.pointer()
 
-        for rel_match in self._env.finditer(expr, match.obj):  # type: ignore
+        for rel_match in path.finditer(match.obj):  # type: ignore
             if projection == Projection.FLAT:
                 patch.addap("/-", rel_match.obj)
             elif projection == Projection.ROOT:
diff --git a/tests/test_query_projection.py b/tests/test_query_projection.py
index 19c3781..ccd1f3a 100644
--- a/tests/test_query_projection.py
+++ b/tests/test_query_projection.py
@@ -159,3 +159,11 @@ def test_sparse_array_selection() -> None:
     assert list(it) == [
         {"categories": [{"products": [{"title": "Beanie", "social": {"shares": 7}}]}]}
     ]
+
+
+def test_pre_compiled_select_many() -> None:
+    expr = "$.*"
+    data = [{"a": 1, "b": 1, "c": 1}, {"a": 2, "b": 2, "c": 2}, {"b": 3, "a": 3}]
+    projection = (jsonpath.compile("a"), "c")
+    it = jsonpath.query(expr, data).select(*projection)
+    assert list(it) == [{"a": 1, "c": 1}, {"a": 2, "c": 2}, {"a": 3}]