add inverted_index, split search into 2 func

Author: shmelevik

search_engine/search_engine.py

@@ -1,19 +1,21 @@
 import re
 from collections import defaultdict
+from itertools import chain
+
+
+def preprocess(text):
+    return re.sub(r'[^\w\s]', '', text).lower()
 
 
 def search(docs: list, words: str) -> list:
     """
     Searches for words in the documents and returns their IDs,
     ranked by the total count of word occurrences in the document text.
     """
-    def preprocess(text):
-        return re.sub(r'[^\w\s]', '', text).lower()
 
     result = []
     words = preprocess(words).split()
 
-
     for doc in docs:
         cleaned_text = preprocess(doc['text'])
         word_count = 0
@@ -30,3 +32,22 @@ def preprocess(text):
     result.sort(key=lambda x: x['word_count'], reverse=True)
 
     return [item['id'] for item in result]
+
+
+def get_inverted_index(docs: list) -> dict:
+    """
+    Builds an inverted index from a list of documents.
+
+    An inverted index maps each unique word found in the documents to a list of 
+    document IDs where that word appears.
+    """
+    result = {}
+
+    words_pit = set(chain.from_iterable(
+            (preprocess(dic['text'])).split()
+        for dic in docs))
+
+    for word in words_pit:
+        result[word] = search(docs, word)
+
+    return result

tests/test_search_engine.py

@@ -1,5 +1,5 @@
 import pytest
-from search_engine.search_engine import search
+from search_engine.search_engine import preprocess, search, get_inverted_index
 
 docs = [
     {'id': 1, 'text': 'hello world'},
@@ -117,4 +117,18 @@ def test_search_multiple_words(self):
         ]
 
         result = search(docs, 'shoot at me')
-        assert result == ['doc2', 'doc1']
+        assert result == ['doc2', 'doc1']
+
+    
+    def test_get_inverted_index(self):
+        doc1 = {'id': 'doc1', 'text': 'some text'}
+        doc2 = {'id': 'doc2', 'text': 'some text too'}
+        docs = [doc1, doc2]
+
+        index = {
+            'some': ['doc1', 'doc2'],
+            'text': ['doc1', 'doc2'],
+            'too': ['doc2']
+        }
+        
+        assert get_inverted_index(docs) == index