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]
|