fix(wiki_tools): return list of full responses for prefix_search return_json

prefix_search with SearchParam.return_json=True returned
result['query']['prefixsearch'] (the page-entry list) and still ran it
through the flat-list comprehension, merging entries across queries into a
single flat list and contradicting the docstring's promise of the full API
response per query.

Mirror the semantic_search behavior: the inner helper now returns the full
raw prefixsearch API response, and the outer function skips the flatten when
return_json=True, returning a list of response dicts (one per query, always a
list even for a single query). The default return_json=False behavior (flat
list of page titles) 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

@@ -150,7 +150,7 @@ def __init__(self, **data):
 
 def prefix_search(
     site: mwclient.client.Site, text: Union[str, SearchParam]
-) -> List[str]:
+) -> Union[List[str], List[dict]]:
     """Standard query. Equivalent to the following mediawiki API call
     api.php?action=query&list=prefixsearch&pssearch=Star Wars.
 
@@ -165,15 +165,17 @@ def prefix_search(
 
     Returns
     -------
-    page_list :
-        List of page titles
+    result:
+        With ``return_json=False`` (default): a flat list of page titles. With
+        ``return_json=True``: a list of raw MediaWiki ``prefixsearch`` API response
+        dicts, one per query (always a list, even for a single query).
     """
     if not isinstance(text, SearchParam):
         query = SearchParam(query=text)
     else:
         query = text
 
-    def prefix_search_(single_text):
+    def prefix_search_(single_text) -> Union[List[str], dict]:
         page_list = list()
         result = site.api(
             "query",
@@ -182,15 +184,16 @@ def prefix_search_(single_text):
             pslimit=query.limit,
             format="json",
         )
-        if len(result["query"]["prefixsearch"]) == 0:
+        if query.debug and len(result["query"]["prefixsearch"]) == 0:
+            print("No results")
+        if query.return_json:
+            return result
+
+        for page in result["query"]["prefixsearch"]:
+            title = page["title"]
             if query.debug:
-                print("No results")
-        else:
-            for page in result["query"]["prefixsearch"]:
-                title = page["title"]
-                if query.debug:
-                    print(title)
-                page_list.append(title)
+                print(title)
+            page_list.append(title)
         return page_list
 
     if query.parallel:
@@ -200,11 +203,17 @@ def prefix_search_(single_text):
     else:
         query_results = [prefix_search_(single_text=sq) for sq in query.query]
 
+    if query.return_json:
+        # Each entry of query_results is the raw API response dict for one query.
+        # Do not flatten dicts; always return the list of responses (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)?
     #  The last option would not change the behavior of the function, but would
-    # return page_list  # original return
+    #  return page_list  # original return
 
 
 def semantic_search(

tests/test_wiki_tools.py

@@ -76,3 +76,55 @@ def test_semantic_search_returns_flat_list_of_titles():
 
     # return_json=False (the default) still yields a flat list of page titles
     assert out == ["Item:OSW1", "Item:OSW2"]
+
+
+def _prefixsearch_result(*titles):
+    """Build a minimal MediaWiki ``prefixsearch`` API result dict."""
+    return {
+        "batchcomplete": "",
+        "query": {
+            "prefixsearch": [
+                {"ns": 0, "title": title, "pageid": idx}
+                for idx, title in enumerate(titles, start=1)
+            ]
+        },
+    }
+
+
+def test_prefix_search_return_json_single_query_returns_list_with_full_response():
+    result = _prefixsearch_result("Star Wars", "Star Trek")
+    site = MagicMock()
+    site.api.return_value = result
+
+    out = wt.prefix_search(site, wt.SearchParam(query="Star", return_json=True))
+
+    # A single query returns a list of one full API response dict, not a flattened
+    # list of individual prefixsearch entries
+    assert isinstance(out, list)
+    assert out == [result]
+    assert out[0]["query"]["prefixsearch"][0]["title"] == "Star Wars"
+
+
+def test_prefix_search_return_json_list_returns_list_of_full_responses():
+    result_a = _prefixsearch_result("Star Wars")
+    result_b = _prefixsearch_result("Voyager")
+    site = MagicMock()
+    site.api.side_effect = [result_a, result_b]
+
+    out = wt.prefix_search(
+        site, wt.SearchParam(query=["Star", "Voy"], return_json=True)
+    )
+
+    # One full API response dict per query, not merged/flattened across queries
+    assert out == [result_a, result_b]
+
+
+def test_prefix_search_returns_flat_list_of_titles():
+    result = _prefixsearch_result("Star Wars", "Star Trek")
+    site = MagicMock()
+    site.api.return_value = result
+
+    out = wt.prefix_search(site, "Star")
+
+    # return_json=False (the default) still yields a flat list of page titles
+    assert out == ["Star Wars", "Star Trek"]