fix(wiki_tools): return list of result dicts for semantic_search return_json

semantic_search with SearchParam.return_json=True previously fed the raw
SMW ask result dicts into the flat-list comprehension, so iterating each
dict yielded its keys and the call returned a useless list like ['query']
instead of the result data.

Skip the flatten when return_json=True and return the list of raw result
dicts (one per query, always a list even for a single query). The default
return_json=False behavior (flat list of page-title fulltext strings) is
unchanged.

Add unit tests covering return_json=True for a single query and a list of
queries, plus return_json=False still returning a flat list of titles.

Author: LukasGold

src/osw/wiki_tools.py

@@ -221,8 +221,10 @@ def semantic_search(
 
     Returns
     -------
-    page_list:
-        List of page titles
+    result:
+        With ``return_json=False`` (default): a flat list of page-title fulltext
+        strings. With ``return_json=True``: a list of raw SMW ``ask`` result dicts,
+        one per query (always a list, even for a single query).
     """
     if not isinstance(query, SearchParam):
         query = SearchParam(query=query)
@@ -231,27 +233,26 @@ def semantic_search_(single_query):
         page_list = list()
         single_query += f"|limit={query.limit}"
         result = site.api("ask", query=single_query, format="json")
-        if len(result["query"]["results"]) == 0:
-            if query.debug:
+        if query.debug:
+            if len(result["query"]["results"]) == 0:
                 print("Query '{}' returned no results".format(single_query))
-        else:
-            if query.debug:
+            else:
                 print(
                     "Query '{}' returned {} results".format(
                         single_query, len(result["query"]["results"])
                     )
                 )
-            if query.return_json:
-                return result
-
-            for page in result["query"]["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)
+        if query.return_json:
+            return result
+
+        for page in result["query"]["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)
         return page_list
 
     if query.parallel:
@@ -261,6 +262,12 @@ def semantic_search_(single_query):
     else:
         query_results = [semantic_search_(single_query=sq) for sq in query.query]
 
+    if query.return_json:
+        # Each entry of query_results is the raw SMW result dict for one query.
+        # Do not flatten dicts; always return the list of result dicts (one per
+        # query), even when only a single query was passed.
+        return query_results
+
     return [item for sublist in query_results for item in sublist]
     # todo: @Simon: a list of lists of strings (sublist for each query in query list)
     #  or a list of strings (results of all queries combined)?

tests/test_wiki_tools.py

@@ -1,3 +1,5 @@
+from unittest.mock import MagicMock
+
 import osw.wiki_tools as wt
 
 
@@ -6,3 +8,71 @@ def test_create_flat_content_structure_from_wikitext():
     expected = [{"Template": {"param": ["value"]}}]
     result = wt.create_flat_content_structure_from_wikitext(wikitext)
     assert result == expected
+
+
+def _ask_result(*titles):
+    """Build a minimal SMW ``ask`` API result dict for the given page titles."""
+    return {
+        "query": {
+            "results": {
+                title: {
+                    "fulltext": title,
+                    "fullurl": f"https://example.org/wiki/{title}",
+                    "namespace": 0,
+                    "exists": "1",
+                    "displaytitle": "",
+                    "printouts": {"SomeProperty": [f"{title}-value"]},
+                }
+                for title in titles
+            }
+        }
+    }
+
+
+def test_semantic_search_return_json_single_query_returns_list_with_dict():
+    result = _ask_result("Item:OSW1")
+    site = MagicMock()
+    site.api.return_value = result
+
+    out = wt.semantic_search(
+        site, wt.SearchParam(query="[[HasType::Category:Item]]", return_json=True)
+    )
+
+    # A single query still returns a list (of one raw result dict), not a bare dict
+    # and not a flattened list of keys
+    assert isinstance(out, list)
+    assert out == [result]
+    # printouts must be preserved intact
+    assert out[0]["query"]["results"]["Item:OSW1"]["printouts"] == {
+        "SomeProperty": ["Item:OSW1-value"]
+    }
+
+
+def test_semantic_search_return_json_list_returns_list_of_dicts():
+    result_a = _ask_result("Item:OSW1")
+    result_b = _ask_result("Item:OSW2")
+    site = MagicMock()
+    site.api.side_effect = [result_a, result_b]
+
+    out = wt.semantic_search(
+        site,
+        wt.SearchParam(
+            query=["[[HasType::Category:Item]]", "[[HasType::Category:Person]]"],
+            return_json=True,
+        ),
+    )
+
+    # A list of queries must return a list of result dicts, one per query
+    assert isinstance(out, list)
+    assert out == [result_a, result_b]
+
+
+def test_semantic_search_returns_flat_list_of_titles():
+    result = _ask_result("Item:OSW1", "Item:OSW2")
+    site = MagicMock()
+    site.api.return_value = result
+
+    out = wt.semantic_search(site, "[[HasType::Category:Item]]")
+
+    # return_json=False (the default) still yields a flat list of page titles
+    assert out == ["Item:OSW1", "Item:OSW2"]