fix: handle empty SMW ask results in semantic_search

Merge pull request #150 from OpenSemanticLab/fix/semantic-search-empty-results

Author: LukasGold

src/osw/wiki_tools.py

@@ -1,4 +1,5 @@
 import getpass
+import warnings
 from typing import Dict, List, Optional, Tuple, Union
 
 import mwclient
@@ -216,6 +217,33 @@ def prefix_search_(single_text) -> Union[List[str], dict]:
     #  return page_list  # original return
 
 
+def _ask_results_as_dict(results: Union[dict, list]) -> dict:
+    """Normalise the ``results`` payload of an SMW ``ask`` response to a mapping.
+
+    SMW serialises a non-empty result set as a JSON object keyed by page title,
+    but an empty one as a JSON array, which would otherwise break attribute
+    access on the result.
+
+    Parameters
+    ----------
+    results :
+        The ``result["query"]["results"]`` payload of an SMW ``ask`` response.
+
+    Returns
+    -------
+    result:
+        The payload normalised to a dict, keyed by page title (or index if a
+        title is unavailable).
+    """
+    if isinstance(results, dict):
+        return results
+    if not results:
+        return {}
+    return {
+        page.get("fulltext", str(index)): page for index, page in enumerate(results)
+    }
+
+
 def semantic_search(
     site: mwclient.client.Site, query: Union[str, List[str], SearchParam]
 ) -> Union[List[str], List[dict]]:
@@ -242,26 +270,39 @@ def semantic_search_(single_query):
         page_list = list()
         single_query += f"|limit={query.limit}"
         result = site.api("ask", query=single_query, format="json")
+        results = _ask_results_as_dict(result["query"]["results"])
+        n = len(results)
         if query.debug:
-            if len(result["query"]["results"]) == 0:
+            if n == 0:
                 print(f"Query '{single_query}' returned no results")
             else:
-                print(
-                    "Query '{}' returned {} results".format(
-                        single_query, len(result["query"]["results"])
-                    )
-                )
+                print(f"Query '{single_query}' returned {n} results")
+        if n >= query.limit:
+            warnings.warn(
+                f"Query '{single_query}' returned {n} results, which meets the "
+                f"requested limit of {query.limit}. Results are truncated - raise "
+                f"the limit or page through with '|offset=' to retrieve the "
+                f"remainder."
+            )
         if query.return_json:
             return result
 
-        for page in result["query"]["results"].values():
+        dropped = 0
+        for page in results.values():
             title = page["fulltext"]
             exists = page["exists"]
             if "#" not in title and query.debug:
                 print(title)
                 # original position of "page_list.append(title)" line
             if exists == "1":
                 page_list.append(title)
+            else:
+                dropped += 1
+        if dropped > 0:
+            warnings.warn(
+                f"Query '{single_query}': {dropped} of {n} results were dropped "
+                f"because the wiki reported them as non-existing pages."
+            )
         return page_list
 
     if query.parallel:

tests/test_wiki_tools.py

@@ -1,5 +1,8 @@
+import warnings
 from unittest.mock import MagicMock
 
+import pytest
+
 import osw.wiki_tools as wt
 
 
@@ -29,6 +32,12 @@ def _ask_result(*titles):
     }
 
 
+def _ask_result_empty():
+    """Build the SMW ``ask`` API result dict for a zero-result query, mirroring
+    SMW's behaviour of serialising an empty result set as a JSON array."""
+    return {"query": {"results": []}}
+
+
 def test_semantic_search_return_json_single_query_returns_list_with_dict():
     result = _ask_result("Item:OSW1")
     site = MagicMock()
@@ -78,6 +87,126 @@ def test_semantic_search_returns_flat_list_of_titles():
     assert out == ["Item:OSW1", "Item:OSW2"]
 
 
+def test_semantic_search_zero_results_returns_empty_list():
+    site = MagicMock()
+    site.api.return_value = _ask_result_empty()
+
+    out = wt.semantic_search(site, "[[HasType::Category:Nonexistent]]")
+
+    assert out == []
+
+
+def test_semantic_search_zero_results_return_json_returns_list_with_dict():
+    result = _ask_result_empty()
+    site = MagicMock()
+    site.api.return_value = result
+
+    out = wt.semantic_search(
+        site,
+        wt.SearchParam(query="[[HasType::Category:Nonexistent]]", return_json=True),
+    )
+
+    assert out == [result]
+
+
+def test_semantic_search_batch_with_one_zero_result_query():
+    result_a = _ask_result("Item:OSW1")
+    result_b = _ask_result_empty()
+    result_c = _ask_result("Item:OSW3")
+    site = MagicMock()
+    site.api.side_effect = [result_a, result_b, result_c]
+
+    out = wt.semantic_search(
+        site,
+        wt.SearchParam(
+            query=[
+                "[[HasType::Category:A]]",
+                "[[HasType::Category:B]]",
+                "[[HasType::Category:C]]",
+            ]
+        ),
+    )
+
+    assert out == ["Item:OSW1", "Item:OSW3"]
+
+
+def test_semantic_search_parallel_batch_with_one_zero_result_query():
+    results = [
+        _ask_result("Item:OSW1"),
+        _ask_result("Item:OSW2"),
+        _ask_result_empty(),
+        _ask_result("Item:OSW4"),
+        _ask_result("Item:OSW5"),
+        _ask_result("Item:OSW6"),
+    ]
+    site = MagicMock()
+    site.api.side_effect = results
+
+    out = wt.semantic_search(
+        site,
+        wt.SearchParam(
+            query=[
+                "[[HasType::Category:A]]",
+                "[[HasType::Category:B]]",
+                "[[HasType::Category:C]]",
+                "[[HasType::Category:D]]",
+                "[[HasType::Category:E]]",
+                "[[HasType::Category:F]]",
+            ]
+        ),
+    )
+
+    assert sorted(out) == [
+        "Item:OSW1",
+        "Item:OSW2",
+        "Item:OSW4",
+        "Item:OSW5",
+        "Item:OSW6",
+    ]
+
+
+def test_semantic_search_truncation_warning():
+    titles = [f"Item:OSW{i}" for i in range(5)]
+    result = _ask_result(*titles)
+    site = MagicMock()
+    site.api.return_value = result
+
+    with pytest.warns(UserWarning, match="truncated"):
+        out = wt.semantic_search(
+            site, wt.SearchParam(query="[[HasType::Category:Item]]", limit=5)
+        )
+
+    assert sorted(out) == sorted(titles)
+
+
+def test_semantic_search_no_truncation_warning_below_limit():
+    titles = [f"Item:OSW{i}" for i in range(5)]
+    result = _ask_result(*titles)
+    site = MagicMock()
+    site.api.return_value = result
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        out = wt.semantic_search(
+            site, wt.SearchParam(query="[[HasType::Category:Item]]", limit=1000)
+        )
+
+    assert not any("truncated" in str(w.message) for w in caught)
+    assert sorted(out) == sorted(titles)
+
+
+def test_semantic_search_exists_drop_warning():
+    result = _ask_result("Item:OSW1", "Item:OSW2")
+    result["query"]["results"]["Item:OSW2"]["exists"] = ""
+    site = MagicMock()
+    site.api.return_value = result
+
+    with pytest.warns(UserWarning, match="non-existing"):
+        out = wt.semantic_search(site, "[[HasType::Category:Item]]")
+
+    assert out == ["Item:OSW1"]
+
+
 def _prefixsearch_result(*titles):
     """Build a minimal MediaWiki ``prefixsearch`` API result dict."""
     return {