Skip to content

Commit

Permalink
Added support for slicing multiple times in Search class (#1771) (#1825)
Browse files Browse the repository at this point in the history
* Added support for slicing multiple times in Search class

Fixes #799

* simplify slicing logic

* A few more slicing unit tests

* Removed unnecessary comment

* simplified slicing logic even more

* added suggested changes

* add more slicing examples to documentation

(cherry picked from commit fb57759)

Co-authored-by: Miguel Grinberg <miguel.grinberg@gmail.com>
  • Loading branch information
github-actions[bot] and miguelgrinberg committed May 17, 2024
1 parent ba57141 commit 1ea2807
Show file tree
Hide file tree
Showing 4 changed files with 79 additions and 18 deletions.
17 changes: 13 additions & 4 deletions docs/search_dsl.rst
Original file line number Diff line number Diff line change
Expand Up @@ -410,16 +410,25 @@ To specify the from/size parameters, use the Python slicing API:

.. code:: python
s = s[10:20]
# {"from": 10, "size": 10}
s = s[10:20]
# {"from": 10, "size": 10}
s = s[:20]
# {"size": 20}
s = s[10:]
# {"from": 10}
s = s[10:20][2:]
# {"from": 12, "size": 8}
If you want to access all the documents matched by your query you can use the
``scan`` method which uses the scan/scroll elasticsearch API:

.. code:: python
for hit in s.scan():
print(hit.title)
for hit in s.scan():
print(hit.title)
Note that in this case the results won't be sorted.

Expand Down
34 changes: 24 additions & 10 deletions elasticsearch_dsl/search_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -352,20 +352,34 @@ def __getitem__(self, n):
# If negative slicing, abort.
if n.start and n.start < 0 or n.stop and n.stop < 0:
raise ValueError("Search does not support negative slicing.")
# Elasticsearch won't get all results so we default to size: 10 if
# stop not given.
s._extra["from"] = n.start or 0
s._extra["size"] = max(
0, n.stop - (n.start or 0) if n.stop is not None else 10
)
return s
slice_start = n.start
slice_stop = n.stop
else: # This is an index lookup, equivalent to slicing by [n:n+1].
# If negative index, abort.
if n < 0:
raise ValueError("Search does not support negative indexing.")
s._extra["from"] = n
s._extra["size"] = 1
return s
slice_start = n
slice_stop = n + 1

old_from = s._extra.get("from")
old_to = None
if "size" in s._extra:
old_to = (old_from or 0) + s._extra["size"]

new_from = old_from
if slice_start is not None:
new_from = (old_from or 0) + slice_start
new_to = old_to
if slice_stop is not None:
new_to = (old_from or 0) + slice_stop
if old_to is not None and old_to < new_to:
new_to = old_to

if new_from is not None:
s._extra["from"] = new_from
if new_to is not None:
s._extra["size"] = max(0, new_to - (new_from or 0))
return s

@classmethod
def from_dict(cls, d):
Expand Down
23 changes: 21 additions & 2 deletions tests/_async/test_search.py
Original file line number Diff line number Diff line change
Expand Up @@ -363,15 +363,34 @@ def test_collapse():
def test_slice():
s = AsyncSearch()
assert {"from": 3, "size": 7} == s[3:10].to_dict()
assert {"from": 0, "size": 5} == s[:5].to_dict()
assert {"from": 3, "size": 10} == s[3:].to_dict()
assert {"size": 5} == s[:5].to_dict()
assert {"from": 3} == s[3:].to_dict()
assert {"from": 0, "size": 0} == s[0:0].to_dict()
assert {"from": 20, "size": 0} == s[20:0].to_dict()
assert {"from": 10, "size": 5} == s[10:][:5].to_dict()
assert {"from": 10, "size": 0} == s[:5][10:].to_dict()
assert {"size": 10} == s[:10][:40].to_dict()
assert {"size": 10} == s[:40][:10].to_dict()
assert {"size": 40} == s[:40][:80].to_dict()
assert {"from": 12, "size": 0} == s[:5][10:][2:].to_dict()
assert {"from": 15, "size": 0} == s[10:][:5][5:].to_dict()
assert {} == s[:].to_dict()
with raises(ValueError):
s[-1:]
with raises(ValueError):
s[4:-1]
with raises(ValueError):
s[-3:-2]


def test_index():
s = AsyncSearch()
assert {"from": 3, "size": 1} == s[3].to_dict()
assert {"from": 3, "size": 1} == s[3][0].to_dict()
assert {"from": 8, "size": 0} == s[3][5].to_dict()
assert {"from": 4, "size": 1} == s[3:10][1].to_dict()
with raises(ValueError):
s[-3]


def test_search_to_dict():
Expand Down
23 changes: 21 additions & 2 deletions tests/_sync/test_search.py
Original file line number Diff line number Diff line change
Expand Up @@ -363,15 +363,34 @@ def test_collapse():
def test_slice():
s = Search()
assert {"from": 3, "size": 7} == s[3:10].to_dict()
assert {"from": 0, "size": 5} == s[:5].to_dict()
assert {"from": 3, "size": 10} == s[3:].to_dict()
assert {"size": 5} == s[:5].to_dict()
assert {"from": 3} == s[3:].to_dict()
assert {"from": 0, "size": 0} == s[0:0].to_dict()
assert {"from": 20, "size": 0} == s[20:0].to_dict()
assert {"from": 10, "size": 5} == s[10:][:5].to_dict()
assert {"from": 10, "size": 0} == s[:5][10:].to_dict()
assert {"size": 10} == s[:10][:40].to_dict()
assert {"size": 10} == s[:40][:10].to_dict()
assert {"size": 40} == s[:40][:80].to_dict()
assert {"from": 12, "size": 0} == s[:5][10:][2:].to_dict()
assert {"from": 15, "size": 0} == s[10:][:5][5:].to_dict()
assert {} == s[:].to_dict()
with raises(ValueError):
s[-1:]
with raises(ValueError):
s[4:-1]
with raises(ValueError):
s[-3:-2]


def test_index():
s = Search()
assert {"from": 3, "size": 1} == s[3].to_dict()
assert {"from": 3, "size": 1} == s[3][0].to_dict()
assert {"from": 8, "size": 0} == s[3][5].to_dict()
assert {"from": 4, "size": 1} == s[3:10][1].to_dict()
with raises(ValueError):
s[-3]


def test_search_to_dict():
Expand Down

0 comments on commit 1ea2807

Please sign in to comment.