diff --git a/elasticsearch_dsl/search_base.py b/elasticsearch_dsl/search_base.py index 5680778c..7a940d4a 100644 --- a/elasticsearch_dsl/search_base.py +++ b/elasticsearch_dsl/search_base.py @@ -743,7 +743,7 @@ def highlight(self, *fields, **kwargs): s._highlight[f] = kwargs return s - def suggest(self, name, text, **kwargs): + def suggest(self, name, text=None, regex=None, **kwargs): """ Add a suggestions request to the search. @@ -754,9 +754,27 @@ def suggest(self, name, text, **kwargs): s = Search() s = s.suggest('suggestion-1', 'Elasticsearch', term={'field': 'body'}) + + # regex query for Completion Suggester + s = Search() + s = s.suggest('suggestion-1', regex='py[thon|py]', completion={'field': 'body'}) """ + if text is None and regex is None: + raise ValueError('You have to pass "text" or "regex" argument.') + if text and regex: + raise ValueError('You can only pass either "text" or "regex" argument.') + if regex and "completion" not in kwargs: + raise ValueError( + '"regex" argument must be passed with "completion" keyword argument.' + ) + s = self._clone() - s._suggest[name] = {"text": text} + if regex: + s._suggest[name] = {"regex": regex} + elif "completion" in kwargs: + s._suggest[name] = {"prefix": text} + else: + s._suggest[name] = {"text": text} s._suggest[name].update(kwargs) return s diff --git a/tests/_async/test_search.py b/tests/_async/test_search.py index 1fc1ac98..6aedb0ed 100644 --- a/tests/_async/test_search.py +++ b/tests/_async/test_search.py @@ -718,3 +718,43 @@ async def test_empty_search(): assert [hit async for hit in s] == [] assert [hit async for hit in s.scan()] == [] await s.delete() # should not error + + +def test_suggest_completion(): + s = AsyncSearch() + s = s.suggest("my_suggestion", "pyhton", completion={"field": "title"}) + + assert { + "suggest": { + "my_suggestion": {"completion": {"field": "title"}, "prefix": "pyhton"} + } + } == s.to_dict() + + +def test_suggest_regex_query(): + s = AsyncSearch() + s = s.suggest("my_suggestion", regex="py[thon|py]", completion={"field": "title"}) + + assert { + "suggest": { + "my_suggestion": {"completion": {"field": "title"}, "regex": "py[thon|py]"} + } + } == s.to_dict() + + +def test_suggest_must_pass_text_or_regex(): + s = AsyncSearch() + with raises(ValueError): + s.suggest("my_suggestion") + + +def test_suggest_can_only_pass_text_or_regex(): + s = AsyncSearch() + with raises(ValueError): + s.suggest("my_suggestion", text="python", regex="py[hton|py]") + + +def test_suggest_regex_must_be_wtih_completion(): + s = AsyncSearch() + with raises(ValueError): + s.suggest("my_suggestion", regex="py[thon|py]") diff --git a/tests/_sync/test_search.py b/tests/_sync/test_search.py index 9a32ee13..dcf0fe6f 100644 --- a/tests/_sync/test_search.py +++ b/tests/_sync/test_search.py @@ -716,3 +716,43 @@ def test_empty_search(): assert [hit for hit in s] == [] assert [hit for hit in s.scan()] == [] s.delete() # should not error + + +def test_suggest_completion(): + s = Search() + s = s.suggest("my_suggestion", "pyhton", completion={"field": "title"}) + + assert { + "suggest": { + "my_suggestion": {"completion": {"field": "title"}, "prefix": "pyhton"} + } + } == s.to_dict() + + +def test_suggest_regex_query(): + s = Search() + s = s.suggest("my_suggestion", regex="py[thon|py]", completion={"field": "title"}) + + assert { + "suggest": { + "my_suggestion": {"completion": {"field": "title"}, "regex": "py[thon|py]"} + } + } == s.to_dict() + + +def test_suggest_must_pass_text_or_regex(): + s = Search() + with raises(ValueError): + s.suggest("my_suggestion") + + +def test_suggest_can_only_pass_text_or_regex(): + s = Search() + with raises(ValueError): + s.suggest("my_suggestion", text="python", regex="py[hton|py]") + + +def test_suggest_regex_must_be_wtih_completion(): + s = Search() + with raises(ValueError): + s.suggest("my_suggestion", regex="py[thon|py]")