Fix bug in builder metadata whitelisting

yeraydiazdiaz · May 15, 2018 · 7127d3d · 7127d3d
1 parent 6e7f218
commit 7127d3d
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 11 deletions.
diff --git a/lunr/builder.py b/lunr/builder.py
@@ -119,27 +119,28 @@ def add(self, doc):
 
             for term in terms:
                 # TODO: term is a Token, should we allow Tokens as keys?
-                term = str(term)
+                term_key = str(term)
 
-                field_terms[term] += 1
-                if term not in self.inverted_index:
+                field_terms[term_key] += 1
+                if term_key not in self.inverted_index:
                     posting = {_field_name: {} for _field_name in self._fields}
                     posting['_index'] = self.term_index
                     self.term_index += 1
-                    self.inverted_index[term] = posting
+                    self.inverted_index[term_key] = posting
 
-                if doc_ref not in self.inverted_index[term][field_name]:
-                    self.inverted_index[term][field_name][doc_ref] = {}
+                if doc_ref not in self.inverted_index[term_key][field_name]:
+                    self.inverted_index[term_key][field_name][doc_ref] = {}
 
                 for metadata_key in self.metadata_whitelist:
                     metadata = term.metadata[metadata_key]
 
-                    if metadata_key not in self.inverted_index[term][doc_ref]:
-                        self.inverted_index[term][field_name][doc_ref][
-                            metadata_key] = []
+                    term_entry = (
+                        self.inverted_index[term_key][field_name][doc_ref])
+                    if metadata_key not in term_entry:
+                        term_entry[metadata_key] = []
 
-                    self.inverted_index[term][field_name][doc_ref][
-                        metadata_key].push(metadata)
+                    self.inverted_index[term_key][field_name][doc_ref][
+                        metadata_key].append(metadata)
 
     def build(self):
         """Builds the index, creating an instance of `lunr.Index`.

diff --git a/tests/test_builder.py b/tests/test_builder.py
@@ -52,6 +52,25 @@ def test_index_is_returned(self):
         assert isinstance(self.index, Index)
 
 
+class TestBuilderAdd:
+
+    def test_builder_metadata_whitelist_includes_metadata_in_index(self):
+        self.builder = Builder()
+        self.builder.ref('id')
+        self.builder.field('title')
+        self.builder.metadata_whitelist = ['position']
+
+        self.builder.add(dict(id='a', title='test', body='missing'))
+        self.builder.add(dict(id='b', title='another test', body='missing'))
+
+        self.builder.inverted_index['test']['title']['a'] == {
+            'position': [[0, 4]]
+        }
+        self.builder.inverted_index['test']['title']['b'] == {
+            'position': [[8, 12]]
+        }
+
+
 class TestBuilderUse:
 
     def setup_method(self, method):