diff --git a/elasticsearch_dsl/search.py b/elasticsearch_dsl/search.py index e5e0b8e26..5273d1b33 100644 --- a/elasticsearch_dsl/search.py +++ b/elasticsearch_dsl/search.py @@ -605,21 +605,39 @@ 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. :arg name: name of the suggestion :arg text: text to suggest on + :arg regex: regex query for Completion Suggester All keyword arguments will be added to the suggestions body. For example:: 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} - s._suggest[name].update(kwargs) + d = s._suggest[name] = kwargs + if regex: + d['regex'] = regex + elif 'completion' in kwargs: + d['prefix'] = text + else: + d['text'] = text return s def to_dict(self, count=False, **kwargs): diff --git a/test_elasticsearch_dsl/test_search.py b/test_elasticsearch_dsl/test_search.py index 53d3c95d6..18817dc42 100644 --- a/test_elasticsearch_dsl/test_search.py +++ b/test_elasticsearch_dsl/test_search.py @@ -505,6 +505,53 @@ def test_suggest(): } } == s.to_dict() + +def test_suggest_completion(): + s = search.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.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.Search() + with raises(ValueError): + s.suggest('my_suggestion') + + +def test_suggest_can_only_pass_text_or_regex(): + s = search.Search() + with raises(ValueError): + s.suggest('my_suggestion', text='python', regex='py[hton|py]') + + +def test_suggest_regex_must_be_wtih_completion(): + s = search.Search() + with raises(ValueError): + s.suggest('my_suggestion', regex='py[thon|py]') + + def test_exclude(): s = search.Search() s = s.exclude('match', title='python')