Redfeine take to leave the receiver in a useful state.

Author: jg-rp

docs/query.md

@@ -45,11 +45,11 @@ print(values[1])
 
 The following `Query` methods all return `self` (the same `Query` instance), so method calls can be chained to further manipulate the underlying iterator.
 
-| Method          | Aliases                 | Description                                        |
-| --------------- | ----------------------- | -------------------------------------------------- |
-| `skip(n: int)`  | `drop`                  | Drop up to _n_ matches from the iterator.          |
-| `limit(n: int)` | `head`, `take`, `first` | Yield at most _n_ matches from the iterator.       |
-| `tail(n: int)`  | `last`                  | Drop matches from the iterator up to the last _n_. |
+| Method          | Aliases         | Description                                        |
+| --------------- | --------------- | -------------------------------------------------- |
+| `skip(n: int)`  | `drop`          | Drop up to _n_ matches from the iterator.          |
+| `limit(n: int)` | `head`, `first` | Yield at most _n_ matches from the iterator.       |
+| `tail(n: int)`  | `last`          | Drop matches from the iterator up to the last _n_. |
 
 ## Terminal methods
 
@@ -64,6 +64,22 @@ These are terminal methods of the `Query` class. They can not be chained.
 | `first_one()` | `one`   | Return the first `JSONPathMatch`, or `None` if there were no matches.                       |
 | `last_one()`  |         | Return the last `JSONPathMatch`, or `None` if there were no matches.                        |
 
+## Take
+
+[`Query.take(self, n: int)`](api.md#jsonpath.Query.take) returns a new `Query` instance, iterating over the next _n_ matches. It leaves the existing query in a safe state, ready to resume iteration of remaining matches.
+
+```python
+from jsonpath import query
+
+it = query("$.some.*", {"some": [0, 1, 2, 3]})
+
+for match in it.take(2):
+    print(match.value)  # 0, 1
+
+for value in it.values():
+    print(value)  # 2, 3
+```
+
 ## Tee
 
 And finally there's `tee()`, which creates multiple independent queries from one query iterator. It is not safe to use the initial `Query` instance after calling `tee()`.

jsonpath/fluent_api.py

@@ -6,6 +6,7 @@
 from typing import TYPE_CHECKING
 from typing import Iterable
 from typing import Iterator
+from typing import List
 from typing import Optional
 from typing import Tuple
 
@@ -33,47 +34,37 @@ def __init__(self, it: Iterable[JSONPathMatch]) -> None:
     def __iter__(self) -> Iterator[JSONPathMatch]:
         return self._it
 
-    def take(self, n: int) -> Query:
+    def limit(self, n: int) -> Query:
         """Limit the query iterator to at most _n_ matches.
 
         Raises:
             ValueError: If _n_ < 0.
         """
         if n < 0:
-            raise ValueError("can't take a negative number of matches")
+            raise ValueError("can't limit by a negative number of matches")
 
         self._it = itertools.islice(self._it, n)
         return self
 
-    def limit(self, n: int) -> Query:
-        """Limit the query iterator to at most _n_ matches.
-
-        `limit()` is an alias of `take()`.
-
-        Raises:
-            ValueError: If _n_ < 0.
-        """
-        return self.take(n)
-
     def head(self, n: int) -> Query:
         """Limit the query iterator to at most the first _n_ matches.
 
-        `head()` is an alias for `take()`.
+        `head()` is an alias for `limit()`.
 
         Raises:
             ValueError: If _n_ < 0.
         """
-        return self.take(n)
+        return self.limit(n)
 
     def first(self, n: int) -> Query:
         """Limit the query iterator to at most the first _n_ matches.
 
-        `first()` is an alias for `take()`.
+        `first()` is an alias for `limit()`.
 
         Raises:
             ValueError: If _n_ < 0.
         """
-        return self.take(n)
+        return self.limit(n)
 
     def drop(self, n: int) -> Query:
         """Skip up to _n_ matches from the query iterator.
@@ -162,3 +153,10 @@ 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))
+
+    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)))

tests/test_fluent_api.py

@@ -123,18 +123,10 @@ def test_query_limit_all() -> None:
 
 def test_query_limit_negative() -> None:
     """Test that we get an exception if limit is negative."""
-    with pytest.raises(ValueError, match="can't take a negative number of matches"):
+    with pytest.raises(ValueError, match="can't limit by a negative number of matches"):
         query("$.some.*", {"some": [0, 1, 2, 3]}).limit(-1)
 
 
-def test_query_take() -> None:
-    """Test that we can limit the number of matches with `take`."""
-    it = query("$.some.*", {"some": [0, 1, 2, 3]}).take(2)
-    matches = list(it)
-    assert len(matches) == 2  # noqa: PLR2004
-    assert [m.obj for m in matches] == [0, 1]
-
-
 def test_query_head() -> None:
     """Test that we can limit the number of matches with `head`."""
     it = query("$.some.*", {"some": [0, 1, 2, 3]}).head(2)
@@ -248,3 +240,12 @@ def test_query_pointers() -> None:
     pointers = list(query("$.some.*", {"some": [0, 1, 2, 3]}).pointers())
     assert len(pointers) == 4  # noqa: PLR2004
     assert pointers[0] == JSONPointer("/some/0")
+
+
+def test_query_take() -> None:
+    """Test that we can take matches from a query iterable."""
+    it = query("$.some.*", {"some": [0, 1, 2, 3]})
+    head = list(it.take(2).values())
+    assert len(head) == 2  # noqa: PLR2004
+    assert head == [0, 1]
+    assert list(it.values()) == [2, 3]