File: auto_complete.py

package info (click to toggle)
python-redisearch-py 1.0.0-2
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 2,356 kB
  • sloc: python: 1,689; makefile: 4
file content (145 lines) | stat: -rw-r--r-- 5,122 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
from redis import Redis, ConnectionPool
from six.moves import xrange

from ._util import to_string

class Suggestion(object):
    """
    Represents a single suggestion being sent or returned from the auto complete server
    """
    def __init__(self, string, score=1.0, payload=None):
        self.string = to_string(string)
        self.payload = to_string(payload)
        self.score = score

    def __repr__(self):
        return self.string


class SuggestionParser(object):
    """
    Internal class used to parse results from the `SUGGET` command.
    This needs to consume either 1, 2, or 3 values at a time from
    the return value depending on what objects were requested
    """
    def __init__(self, with_scores, with_payloads, ret):
        self.with_scores = with_scores
        self.with_payloads = with_payloads

        if with_scores and with_payloads:
            self.sugsize = 3
            self._scoreidx = 1
            self._payloadidx = 2
        elif with_scores:
            self.sugsize = 2
            self._scoreidx = 1
        elif with_payloads:
            self.sugsize = 2
            self._payloadidx = 1
        else:
            self.sugsize = 1
            self._scoreidx = -1

        self._sugs = ret

    def __iter__(self):
        for i in xrange(0, len(self._sugs), self.sugsize):
            ss = self._sugs[i]
            score = float(self._sugs[i + self._scoreidx]) if self.with_scores else 1.0
            payload = self._sugs[i + self._payloadidx] if self.with_payloads else None
            yield Suggestion(ss, score, payload)


class AutoCompleter(object):
    """
    A client to RediSearch's AutoCompleter API

    It provides prefix searches with optionally fuzzy matching of prefixes    
    """

    SUGADD_COMMAND = "FT.SUGADD"
    SUGDEL_COMMAND = "FT.SUGDEL"
    SUGLEN_COMMAND = "FT.SUGLEN"
    SUGGET_COMMAND = "FT.SUGGET"

    INCR = 'INCR'
    WITHSCORES = 'WITHSCORES'
    FUZZY = 'FUZZY'
    WITHPAYLOADS = 'WITHPAYLOADS'

    def __init__(self, key, host='localhost', port=6379, conn = None):
        """
        Create a new AutoCompleter client for the given key, and optional host and port

        If conn is not None, we employ an already existing redis connection
        """

        self.key = key
        self.redis = conn if conn is not None else Redis(
            connection_pool = ConnectionPool(host=host, port=port))

    def add_suggestions(self,  *suggestions, **kwargs):
        """
        Add suggestion terms to the AutoCompleter engine. Each suggestion has a score and string.

        If kwargs['increment'] is true and the terms are already in the server's dictionary, we increment their scores
        """
        pipe = self.redis.pipeline()
        for sug in suggestions:
            args = [AutoCompleter.SUGADD_COMMAND, self.key, sug.string, sug.score]
            if kwargs.get('increment'):
                args.append(AutoCompleter.INCR)
            if sug.payload:
                args.append('PAYLOAD')
                args.append(sug.payload)

            pipe.execute_command(*args)

        return pipe.execute()[-1]



    def len(self):
        """
        Return the number of entries in the AutoCompleter index
        """
        return self.redis.execute_command(AutoCompleter.SUGLEN_COMMAND, self.key)

    def delete(self, string):
        """
        Delete a string from the AutoCompleter index.
        Returns 1 if the string was found and deleted, 0 otherwise
        """
        return self.redis.execute_command(AutoCompleter.SUGDEL_COMMAND, self.key, string)

    def get_suggestions(self, prefix, fuzzy = False, num = 10, with_scores = False, with_payloads=False):
        """
        Get a list of suggestions from the AutoCompleter, for a given prefix

        ### Parameters:
        - **prefix**: the prefix we are searching. **Must be valid ascii or utf-8**
        - **fuzzy**: If set to true, the prefix search is done in fuzzy mode. 
            **NOTE**: Running fuzzy searches on short (<3 letters) prefixes can be very slow, and even scan the entire index.
        - **with_scores**: if set to true, we also return the (refactored) score of each suggestion. 
          This is normally not needed, and is NOT the original score inserted into the index
        - **with_payloads**: Return suggestion payloads
        - **num**: The maximum number of results we return. Note that we might return less. The algorithm trims irrelevant suggestions.
        
        Returns a list of Suggestion objects. If with_scores was False, the score of all suggestions is 1.
        """

        args = [AutoCompleter.SUGGET_COMMAND, self.key, prefix, 'MAX', num]
        if fuzzy:
            args.append(AutoCompleter.FUZZY)
        if with_scores:
            args.append(AutoCompleter.WITHSCORES)
        if with_payloads:
            args.append(AutoCompleter.WITHPAYLOADS)

        ret = self.redis.execute_command(*args)
        results = []
        if not ret:
            return results

        parser = SuggestionParser(with_scores, with_payloads, ret)
        return [s for s in parser]