From 657e1c4170437436a0405923355378ac47bf817a Mon Sep 17 00:00:00 2001
From: Joe Paul <joeirimpan@gmail.com>
Date: Tue, 18 Sep 2018 12:51:45 +0530
Subject: [PATCH] feat: Add spellcheck, dict operation commands

---
 redisearch/client.py | 101 +++++++++++++++++++++++++++++++++++++++++++
 test/test.py         |  35 +++++++++++++++
 2 files changed, 136 insertions(+)

diff --git a/redisearch/client.py b/redisearch/client.py
index ef20858..35562e8 100644
--- a/redisearch/client.py
+++ b/redisearch/client.py
@@ -109,6 +109,10 @@ class Client(object):
     DEL_CMD = 'FT.DEL'
     AGGREGATE_CMD = 'FT.AGGREGATE'
     CURSOR_CMD = 'FT.CURSOR'
+    SPELLCHECK_CMD = 'FT.SPELLCHECK'
+    DICT_ADD_CMD = 'FT.DICTADD'
+    DICT_DEL_CMD = 'FT.DICTDEL'
+    DICT_DUMP_CMD = 'FT.DICTDUMP'
 
 
     NOOFFSETS = 'NOOFFSETS'
@@ -383,3 +387,100 @@ def aggregate(self, query):
 
         res = AggregateResult(rows, cursor, schema)
         return res
+
+    def spellcheck(self, query, distance=None, include=None, exclude=None):
+        """
+        Issue a spellcheck query
+
+        ### Parameters
+
+        **query**: search query.
+        **distance***: the maximal Levenshtein distance for spelling suggestions (default: 1, max: 4).
+        **include**: specifies an inclusion custom dictionary.
+        **exclude**: specifies an exclusion custom dictionary.
+        """
+        cmd = [self.SPELLCHECK_CMD, self.index_name, query]
+        if distance:
+            cmd.extend(['DISTANCE', distance])
+
+        if include:
+            cmd.extend(['TERMS', 'INCLUDE', include])
+
+        if exclude:
+            cmd.extend(['TERMS', 'EXCLUDE', exclude])
+
+        raw = self.redis.execute_command(*cmd)
+
+        corrections = {}
+        if raw == 0:
+            return corrections
+
+        for _correction in raw:
+            if isinstance(_correction, long) and _correction == 0:
+                continue
+
+            if len(_correction) != 3:
+                continue
+            if not _correction[2]:
+                continue
+            if not _correction[2][0]:
+                continue
+
+            # For spellcheck output
+            # 1)  1) "TERM"
+            #     2) "{term1}"
+            #     3)  1)  1)  "{score1}"
+            #             2)  "{suggestion1}"
+            #         2)  1)  "{score2}"
+            #             2)  "{suggestion2}"
+            #
+            # Following dictionary will be made
+            # corrections = {
+            #     '{term1}': [
+            #         {'score': '{score1}', 'suggestion': '{suggestion1}'},
+            #         {'score': '{score2}', 'suggestion': '{suggestion2}'}
+            #     ]
+            # }
+            corrections[_correction[1]] = [
+                {'score': _item[0], 'suggestion':_item[1]}
+                for _item in _correction[2]
+            ]
+
+        return corrections
+
+    def dict_add(self, name, *terms):
+        """Adds terms to a dictionary.
+
+        ### Parameters
+
+        - **name**: Dictionary name.
+        - **terms**: List of items for adding to the dictionary.
+        """
+        cmd = [self.DICT_ADD_CMD, name]
+        cmd.extend(terms)
+        raw = self.redis.execute_command(*cmd)
+        return raw
+
+    def dict_del(self, name, *terms):
+        """Deletes terms from a dictionary.
+
+        ### Parameters
+
+        - **name**: Dictionary name.
+        - **terms**: List of items for removing from the dictionary.
+        """
+        cmd = [self.DICT_DEL_CMD, name]
+        cmd.extend(terms)
+        raw = self.redis.execute_command(*cmd)
+        return raw
+
+    def dict_dump(self, name):
+        """Dumps all terms in the given dictionary.
+
+        ### Parameters
+
+        - **name**: Dictionary name.
+        """
+        cmd = [self.DICT_DUMP_CMD, name]
+        raw = self.redis.execute_command(*cmd)
+        return raw
diff --git a/test/test.py b/test/test.py
index a84086c..ee1b77b 100644
--- a/test/test.py
+++ b/test/test.py
@@ -534,6 +534,41 @@ def testAlterSchemaAdd(self):
             res = client.search(q)
             self.assertEqual(1, res.total)
 
+    def testSpellCheck(self):
+        client = self.getCleanClient('idx')
+        client.create_index((TextField('f1'), TextField('f2')))
+
+        client.add_document('doc1', f1='some valid content', f2='this is sample text')
+        client.add_document('doc2', f1='very important', f2='lorem ipsum')
+
+        for i in self.retry_with_reload():
+            res = client.spellcheck('impornant')
+            self.assertEqual('important', res['impornant'][0]['suggestion'])
+
+            res = client.spellcheck('contnt')
+            self.assertEqual('content', res['contnt'][0]['suggestion'])
+
+    def testDictOps(self):
+        client = self.getCleanClient('idx')
+        client.create_index((TextField('f1'), TextField('f2')))
+
+        for i in self.retry_with_reload():
+            # Add three items
+            res = client.dict_add('custom_dict', 'item1', 'item2', 'item3')
+            self.assertEqual(3, res)
+
+            # Remove one item
+            res = client.dict_del('custom_dict', 'item2')
+            self.assertEqual(1, res)
+
+            # Dump dict and inspect content
+            res = client.dict_dump('custom_dict')
+            self.assertEqual(['item1', 'item3'], res)
+
+            # Remove rest of the items before reload
+            client.dict_del('custom_dict', *res)
+
+
 if __name__ == '__main__':
 
     unittest.main()