add search for multiple words

Author: shmelevik

search_engine/search_engine.py

@@ -1,21 +1,32 @@
 import re
+from collections import defaultdict
 
-def search(docs: list, word: str) -> list:
+
+def search(docs: list, words: str) -> list:
     """
-    Searches for a word in the documents and returns their IDs.
+    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 = []
-    word = preprocess(word)
+    words = preprocess(words).split()
+
 
     for doc in docs:
         cleaned_text = preprocess(doc['text'])
-        word_count = cleaned_text.split().count(word)
+        word_count = 0
+
+        for word in words:
+            word_count += cleaned_text.split().count(word)
 
         if word_count:
-            result.append((doc['id'], word_count))
+            result.append({
+                'id': doc['id'],
+                'word_count': word_count
+            })
+
+    result.sort(key=lambda x: x['word_count'], reverse=True)
 
-    result.sort(key=lambda x: x[1], reverse=True)
-    return [id for id, _ in result]
+    return [item['id'] for item in result]

tests/test_search_engine.py

@@ -94,4 +94,27 @@ def test_search_ranging(self):
         ]
 
         result = search(docs, 'shoot')
+        assert result == ['doc2', 'doc1']
+
+    def test_search_multiple_words(self):
+        """
+        Test that the search function returns documents containing all the words
+        in the search query, and sorts them based on the frequency of the words' occurrence.
+
+        The test ensures the following:
+        - The search can handle multiple words (i.e., 'shoot at me').
+        - Documents are returned in the order of their relevance (based on word frequency).
+        - Words in the search query are treated independently (i.e., not as a single phrase).
+        """
+        doc1 = "I can't shoot straight unless I've had a pint!"
+        doc2 = "Don't shoot shoot shoot that thing at me."
+        doc3 = "I'm your shooter."
+
+        docs = [
+            {'id': 'doc1', 'text': doc1},
+            {'id': 'doc2', 'text': doc2},
+            {'id': 'doc3', 'text': doc3},
+        ]
+
+        result = search(docs, 'shoot at me')
         assert result == ['doc2', 'doc1']